python_scripts/training_and_monitoring_auto.py

# ---------- fine_tune_v15_resume_progress_eta_fixed.py ----------
# (주식회사 엠케이디 | MKD)
# MKDCore Continuous Fine-Tuning Pipeline
# (ETA Dashboard + Resume Fix + Segment Summary + Token Control)

import os, sys, time, gc, pickle, random, datetime, subprocess, math, traceback, signal
import torch
from torch import nn, optim
from tqdm import tqdm
import psutil
import wandb
from mkd_core import MKDCore, MyTokenizer, SEQ_LEN

# ============================================================
# CONFIGURATION
# ============================================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BASE_DIR = "/home/mkd-s003/AI/finetune"

TOKENIZER_JSON = os.path.join(BASE_DIR, "prod_tokenizer_v2.json")
TRAIN_PKL = os.path.join(BASE_DIR, "token_cache_train.pkl")
VAL_PKL   = os.path.join(BASE_DIR, "token_cache_val.pkl")
MODEL_BASE = os.path.join(BASE_DIR, "mkdcore_v2.pt")

CKPT_DIR   = os.path.join(BASE_DIR, "checkpoints_ft")
LOG_DIR    = os.path.join(BASE_DIR, "logs")
FINE_DIR   = os.path.join(BASE_DIR, "finetune_model")
LOG_REASON = os.path.join(LOG_DIR, "reason_of_interruption.log")
PROGRESS_PKL = os.path.join(BASE_DIR, "segment_progress.pkl")
SEGMENT_SUMMARY = os.path.join(LOG_DIR, "segment_summary.txt")

for d in [CKPT_DIR, LOG_DIR, FINE_DIR]:
    os.makedirs(d, exist_ok=True)

MASTER_LOG = os.path.join(LOG_DIR, "training_log.txt")
EVENT_LOG  = os.path.join(LOG_DIR, "training_event_log.txt")

# Segmenting config
MAX_TRAIN_TOKENS_SEG = 20_000_000
MAX_VAL_TOKENS_SEG   = 2_300_000
TOTAL_TOKENS = 4_300_000_000
TOKENS_PER_SEG = MAX_TRAIN_TOKENS_SEG + MAX_VAL_TOKENS_SEG
TOTAL_SEGS = max(1, math.ceil(TOTAL_TOKENS / TOKENS_PER_SEG))

# Training config
EPOCHS = 16
STEPS_PER_EPOCH = 20_000
LR = 2e-5
LR_DECAY = 0.9
GRAD_CLIP = 1.0
SAVE_EVERY = 5000
RESTART_DELAY = 60


# ============================================================
# LOGGING UTILITIES
# ============================================================
def _ts(): return datetime.datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")

def _append(path, text):
    with open(path, "a", encoding="utf-8") as f:
        f.write(text + "\n")

def log(msg, master=True):
    line = f"{_ts()} {msg}"
    _append(EVENT_LOG if master else MASTER_LOG, line)
    print(line)

def log_both(msg):
    line = f"{_ts()} {msg}"
    _append(EVENT_LOG, line); _append(MASTER_LOG, line)
    print(line)

def log_reason(reason):
    _append(LOG_REASON, f"{_ts()} {reason}")
    log(f"Interruption reason: {reason}")

# ============================================================
# SYSTEM STATUS
# ============================================================
def system_status():
    try:
        gpu_query = subprocess.check_output([
            "nvidia-smi",
            "--query-gpu=temperature.gpu,utilization.gpu,memory.used,memory.total",
            "--format=csv,noheader,nounits"
        ]).decode().strip().split(", ")
        temp, util, mem_used, mem_total = map(str.strip, gpu_query)
        used_gb = float(mem_used) / 1024
        total_gb = float(mem_total) / 1024
        gpu_status = f"GPU={used_gb:.2f}/{total_gb:.0f}GB | {temp}°C | Util {util}%"
    except Exception:
        gpu_status = "GPU=?GB | ?°C | Util ?%"
    cpu = psutil.cpu_percent(interval=0.1)
    ram = psutil.virtual_memory().percent
    return f"{gpu_status} | CPU {cpu:.0f}% | RAM {ram:.1f}%"

# ============================================================
# SEGMENT SUMMARY WRITER
# ============================================================
def write_segment_summary(seg_idx, train_start, val_start):
    done = train_start + val_start
    remaining = max(0, TOTAL_TOKENS - done)
    header = (
        f"\n{_ts()}\n"
        f"================ SEGMENT {seg_idx} SUMMARY ================\n"
        f"Train Start: {train_start:,} | Val Start: {val_start:,}\n"
        f"Tokens Done: {done:,} / {TOTAL_TOKENS:,}\n"
        f"Tokens Remaining: {remaining:,}\n"
        f"===========================================================\n"
    )
    # Write to segment_summary.txt
    _append(SEGMENT_SUMMARY, header)
    # Also to event log and stdout (via log())
    log(header.strip())

# ============================================================
# SAVE / LOAD PROGRESS
# ============================================================
def save_progress(train_pos, val_pos, segment_completed=False):
    with open(PROGRESS_PKL, "wb") as f:
        pickle.dump({
            "train_pos": int(train_pos),
            "val_pos": int(val_pos),
            "segment_completed": segment_completed
        }, f)

def load_progress():
    if not os.path.exists(PROGRESS_PKL):
        return {"train_pos": 0, "val_pos": 0, "segment_completed": False}
    with open(PROGRESS_PKL, "rb") as f:
        return pickle.load(f)

# ============================================================
# INTERRUPT HANDLER
# ============================================================
def ctrl_c_handler_factory(model, opt, scaler, scheduler, segment_idx,
                           current_epoch_ref, current_step_ref, train_pos, val_pos):
    def handle_ctrl_c(sig, frame):
        epoch = current_epoch_ref["val"]; step = current_step_ref["val"]
        pct = ((step + 1) / STEPS_PER_EPOCH) * 100
        log_both(f"Manual stop detected → Segment {segment_idx}, Epoch {epoch + 1}, Step {step + 1}/{STEPS_PER_EPOCH} ({pct:.1f}%)")
        ckpt_dir = os.path.join(CKPT_DIR, f"segment_{segment_idx:03d}")
        os.makedirs(ckpt_dir, exist_ok=True)
        ckpt_path = os.path.join(ckpt_dir, f"interrupt_ep{epoch:02d}_step{step:05d}.pt")
        torch.save({
            "segment": segment_idx, "epoch": epoch, "step": step,
            "model": model.state_dict(), "optimizer": opt.state_dict(),
            "scheduler": scheduler.state_dict() if scheduler else {},
            "scaler": scaler.state_dict(),
            "train_pos": train_pos, "val_pos": val_pos,
            "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }, ckpt_path)
        log_reason(f"Manual stop → Segment {segment_idx}, Epoch {epoch + 1}, Step {step + 1}/{STEPS_PER_EPOCH}")
        sys.exit(0)
    return handle_ctrl_c

# ============================================================
# CHECKPOINT MANAGEMENT
# ============================================================
def segment_dir(seg):
    p = os.path.join(CKPT_DIR, f"segment_{seg:03d}")
    os.makedirs(p, exist_ok=True)
    return p

def cleanup_old_ckpts(seg):
    p = segment_dir(seg)
    files = sorted([f for f in os.listdir(p) if f.endswith(".pt")],
                   key=lambda x: os.path.getmtime(os.path.join(p, x)))
    if len(files) > 2:
        for f in files[:-2]:
            os.remove(os.path.join(p, f)); log(f"Old checkpoint removed: {f}")

def cleanup_old_finetunes():
    files = sorted(
        [os.path.join(FINE_DIR, f) for f in os.listdir(FINE_DIR) if f.endswith(".pt")],
        key=os.path.getmtime)
    if len(files) > 3:
        for f in files[:-3]:
            os.remove(f); log(f"Old finetune model removed: {os.path.basename(f)}")

def save_ckpt(seg, ep, st, model, opt, sch, sc, tpos, vpos):
    path = os.path.join(segment_dir(seg), f"checkpoint_ep{ep:02d}_step{st:05d}.pt")
    torch.save({
        "segment": seg, "epoch": ep, "step": st,
        "model": model.state_dict(), "optimizer": opt.state_dict(),
        "scheduler": sch.state_dict(), "scaler": sc.state_dict(),
        "train_pos": tpos, "val_pos": vpos
    }, path)
    log_both(f"Checkpoint saved → Segment {seg}, Epoch {ep + 1}, Step {st + 1}/{STEPS_PER_EPOCH}")
    cleanup_old_ckpts(seg)

def load_latest_ckpt(model, opt, sch, sc):
    seg_dirs = [d for d in os.listdir(CKPT_DIR) if d.startswith("segment_")]
    ckpts = [os.path.join(CKPT_DIR, d, f)
             for d in seg_dirs
             for f in os.listdir(os.path.join(CKPT_DIR, d))
             if f.endswith(".pt")]
    if not ckpts: return None, 0, 0
    latest = max(ckpts, key=os.path.getmtime)
    data = torch.load(latest, map_location=DEVICE)
    model.load_state_dict(data["model"], strict=False)
    opt.load_state_dict(data["optimizer"])
    if "scheduler" in data and data["scheduler"]:
        sch.load_state_dict(data["scheduler"])
    sc.load_state_dict(data["scaler"])
    seg, ep, st = data.get("segment", 1), data.get("epoch", 0), data.get("step", 0)
    log_both(f"Resumed training at → Segment {seg}, Epoch {ep + 1}, Step {st + 1}/{STEPS_PER_EPOCH}")
    return data, ep, st

# ============================================================
# DATA LOADING
# ============================================================
def load_tokens_segment(pkl, start, max_tokens, label):
    tokens, loaded, consumed = [], 0, 0
    log(f"Loading {label} | start={start:,} | limit={max_tokens:,}")
    with open(pkl, "rb") as f:
        while loaded < max_tokens:
            try: seq = pickle.load(f)
            except EOFError: break
            if not isinstance(seq, list) or not seq: continue
            L = len(seq)
            if consumed + L <= start: consumed += L; continue
            s = max(0, start - consumed); t = min(L - s, max_tokens - loaded)
            tokens.extend(seq[s:s + t]); loaded += t; consumed += L
            if loaded >= max_tokens: break
    new_offset = start + loaded
    log(f"Loaded {loaded:,} tokens | new_offset={new_offset:,} | {label}")
    return torch.tensor(tokens, dtype=torch.long), new_offset

# ============================================================
# TRAINING UTILITIES
# ============================================================
def get_batch(src, seq_len):
    i = random.randint(0, len(src) - seq_len - 1)
    x = src[i:i+seq_len].unsqueeze(0).to(DEVICE)
    y = src[i+1:i+seq_len+1].unsqueeze(0).to(DEVICE)
    return x, y

@torch.no_grad()
def eval_loss(model, loss_fn, vocab, val):
    model.eval(); total = 0.0
    for _ in range(10):
        x, y = get_batch(val, SEQ_LEN)
        with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()):
            logits = model(x); loss = loss_fn(logits.view(-1, vocab), y.view(-1))
        total += loss.item()
    return total / 10

def fmt_time(seconds):
    seconds = max(0, float(seconds))
    m, s = divmod(int(seconds), 60); h, m = divmod(m, 60)
    return f"{h}h {m:02d}m" if h > 0 else f"{m:02d}m {s:02d}s"

# ============================================================
# MAIN TRAINING FUNCTION
# ============================================================
def train_segment(seg_idx):
    torch.cuda.empty_cache(); gc.collect()
    log(f"Segment {seg_idx}/{TOTAL_SEGS} STARTED")

    tok = MyTokenizer(TOKENIZER_JSON); vocab = tok.vocab_size

    # ✅ Load latest fine-tuned or base model
    model = MKDCore(vocab).to(DEVICE)
    try:
        finetunes = sorted(
            [os.path.join(FINE_DIR, f) for f in os.listdir(FINE_DIR) if f.endswith(".pt")],
            key=os.path.getmtime)
        if finetunes:
            latest = finetunes[-1]
            model.load_state_dict(torch.load(latest, map_location=DEVICE), strict=False)
            log(f"Loaded latest fine-tuned model → {latest}")
        else:
            model.load_state_dict(torch.load(MODEL_BASE, map_location=DEVICE), strict=False)
            log(f"No fine-tuned model found, starting from base → {MODEL_BASE}")
    except Exception as e:
        log(f"⚠️ Could not load fine-tuned model: {e}")

    opt = optim.AdamW(model.parameters(), lr=LR, betas=(0.9,0.95))
    sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=EPOCHS)
    loss_fn = nn.CrossEntropyLoss()
    sc = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())

    state = None
    start_ep, start_st = 0, 0
    prog = load_progress()

    # Resume from checkpoint ONLY if unfinished
    if not prog.get("segment_completed", False):
        state, start_ep, start_st = load_latest_ckpt(model, opt, sch, sc)
        if state is None:
            start_ep, start_st = 0, 0
    else:
        start_ep, start_st = 0, 0
        # clear old segment ckpts so each new segment starts clean
        for d in os.listdir(CKPT_DIR):
            path = os.path.join(CKPT_DIR, d)
            if os.path.isdir(path):
                for f in os.listdir(path):
                    try: os.remove(os.path.join(path, f))
                    except Exception: pass
        log("Cleared old checkpoints for new segment.")

    # ✅ Token progression: move only after successful segment completion
    if prog.get("segment_completed", False):
        train_start = prog["train_pos"] + MAX_TRAIN_TOKENS_SEG
        val_start   = prog["val_pos"]   + MAX_VAL_TOKENS_SEG
        save_progress(train_start, val_start, segment_completed=False)
        log(f"Starting NEW segment block → Train start={train_start:,}, Val start={val_start:,}")
    else:
        train_start = prog["train_pos"]
        val_start   = prog["val_pos"]
        log(f"Resuming same segment (unfinished) → Train start={train_start:,}, Val start={val_start:,}")

    # ✅ Write per-segment banner summary (terminal + logs file + event log)
    write_segment_summary(seg_idx, train_start, val_start)

    train, new_t = load_tokens_segment(TRAIN_PKL, train_start, MAX_TRAIN_TOKENS_SEG, "train")
    val,   new_v = load_tokens_segment(VAL_PKL,   val_start,   MAX_VAL_TOKENS_SEG,   "val")

    wandb.init(
        project="MKD_FineTuning",
        name=f"Segment_{seg_idx}",
        id=f"segment_{seg_idx}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}",
        resume="allow",
        config={
            "segment": seg_idx,
            "epochs": EPOCHS,
            "lr": LR,
            "tokens_per_segment": TOKENS_PER_SEG,
            "total_segments": TOTAL_SEGS
        }
    )

    current_step  = {"val": start_st}
    current_epoch = {"val": start_ep}
    signal.signal(signal.SIGINT,
        ctrl_c_handler_factory(model, opt, sc, sch, seg_idx, current_epoch, current_step, train_start, val_start))

    seg_start_time = time.time()

    for ep in range(start_ep, EPOCHS):
        current_epoch["val"] = ep; model.train(); total_loss = 0.0
        with tqdm(total=STEPS_PER_EPOCH,
                  desc=f"[Seg {seg_idx}/{TOTAL_SEGS} | Rem {TOTAL_SEGS - seg_idx}] Epoch {ep+1}/{EPOCHS}",
                  ncols=170) as bar:
            for st in range(start_st, STEPS_PER_EPOCH):
                current_step["val"] = st
                try:
                    x, y = get_batch(train, SEQ_LEN)
                    with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()):
                        logits = model(x); loss = loss_fn(logits.view(-1, vocab), y.view(-1))
                    sc.scale(loss).backward(); sc.unscale_(opt)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
                    sc.step(opt); sc.update(); opt.zero_grad(set_to_none=True)
                    total_loss += loss.item(); avg = total_loss / (st + 1)

                    elapsed = time.time() - seg_start_time
                    progress = ((ep*STEPS_PER_EPOCH)+(st+1))/(EPOCHS*STEPS_PER_EPOCH)
                    per_seg_total_est = elapsed / max(progress,1e-6)
                    eta_seg = per_seg_total_est - elapsed
                    remaining_segments = max(0, TOTAL_SEGS - seg_idx)
                    eta_total = eta_seg + remaining_segments * per_seg_total_est

                    bar.set_postfix_str(
                        f"Loss={avg:.4f}, ETA(seg)={fmt_time(eta_seg)}, "
                        f"ETA(total)={fmt_time(eta_total)}, {system_status()}"
                    )
                    bar.update(1)

                    if (st + 1) % SAVE_EVERY == 0:
                        save_ckpt(seg_idx, ep, st, model, opt, sch, sc, train_start, val_start)

                except Exception as e:
                    log(f"Training interrupted at Seg {seg_idx}, Ep {ep+1}, St {st+1}: {e}")
                    save_ckpt(seg_idx, ep, st, model, opt, sch, sc, train_start, val_start)
                    log_reason(f"Unexpected stop at Seg {seg_idx}, Ep {ep+1}, St {st+1}")
                    raise e

            val_loss = eval_loss(model, loss_fn, tok.vocab_size, val)
            ppl = math.exp(val_loss)
            log(f"Epoch {ep+1} done | ValLoss {val_loss:.4f} | PPL {ppl:.2f}")
            wandb.log({"val_loss": val_loss, "perplexity": ppl, "epoch": ep+1})
            sch.step()
            for g in opt.param_groups: g["lr"] *= LR_DECAY
            start_st = 0

    wandb.finish()
    model_save = os.path.join(FINE_DIR,
        f"mkdcore_v2_segment_finetuned_{datetime.datetime.now().strftime('%Y%m%d_%H%M')}.pt")
    torch.save(model.state_dict(), model_save)
    log(f"Final model saved → {model_save}")
    cleanup_old_finetunes()
    save_progress(new_t, new_v, segment_completed=True)
    log(f"Segment {seg_idx} completed successfully.")

# ============================================================
# MAIN LOOP
# ============================================================
if __name__ == "__main__":
    seg = 1
    while True:
        try:
            train_segment(seg); seg += 1
        except KeyboardInterrupt:
            log_both("Training stopped manually by user."); sys.exit(0)
        except Exception as e:
            traceback.print_exc()
            log(f"ERROR in seg {seg}: {e}")
            torch.cuda.empty_cache(); gc.collect()
            log(f"Retrying after {RESTART_DELAY}s...")
            time.sleep(RESTART_DELAY)
            continue

        prog = load_progress()
        t, v = prog["train_pos"], prog["val_pos"]

        if seg > TOTAL_SEGS:
            log_both("Training COMPLETED — all segments processed."); break

        try:
            if t >= os.path.getsize(TRAIN_PKL) or v >= os.path.getsize(VAL_PKL):
                log_both("Training COMPLETED — all tokens processed."); break
        except Exception:
            pass

        for s in range(RESTART_DELAY, 0, -1):
            print(f"\rNext segment in {s:02d}s | {system_status()}", end="", flush=True)
            time.sleep(1)
        print()