python_scripts/build_token_cache_mkd_auto_resume_final.py

# ---------- build_token_cache_mkd_auto_resume_v6.py ----------
# (주식회사엠케이디 | MKD)
# Fully automated, resumable, logging-enabled token cache system
# Includes daily logs, auto PKL/TMP rotation, and safe checkpoints

import os
import time
import pickle
import psutil
import subprocess
import sys
import torch
import gc
from datetime import datetime
from tokenizers import Tokenizer

# ============================================================
# CONFIGURATION
# ============================================================
BASE_DIR = "/home/mkd-s003/AI/finetune"
TOKENIZER_PATH = os.path.join(BASE_DIR, "prod_tokenizer_v2.json")
TEXT_PATH = os.path.join(BASE_DIR, "clean_text.txt")

TRAIN_PKL = os.path.join(BASE_DIR, "token_cache_train.pkl")
VAL_PKL = os.path.join(BASE_DIR, "token_cache_val.pkl")
TMP_PKL = os.path.join(BASE_DIR, "token_cache_progress.tmp")
LOG_FILE = os.path.join(BASE_DIR, "token_cache_log.txt")

TOKEN_LIMIT_TRAIN = 200_000_000
TOKEN_LIMIT_VAL   = 20_000_000
SPLIT_RATIO = 0.9
SAVE_EVERY = 500_000
BATCH_LINES = 3_000_000
RESTART_DELAY = 60

# ============================================================
# LOGGING UTILITY
# ============================================================
def log_message(msg):
    """Write timestamped message to both console and log file"""
    stamp = datetime.now().strftime("[%Y-%m-%d %H:%M:%S]")
    text = f"{stamp} {msg}"
    print(text)
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(text + "\n")

# ============================================================
# MONITOR UTILITIES
# ============================================================
def system_status():
    cpu = psutil.cpu_percent(interval=0.2)
    mem = psutil.virtual_memory().percent
    gpu = ""
    if torch.cuda.is_available():
        try:
            util = subprocess.check_output(
                ["nvidia-smi", "--query-gpu=utilization.gpu",
                 "--format=csv,noheader,nounits"]
            ).decode().strip().split("\n")[0]
            gpu = f" | GPU Util {util}%"
        except Exception:
            gpu = ""
    return f"CPU {cpu:.0f}% | RAM {mem:.1f}%{gpu}"

# ============================================================
# BACKUP MANAGEMENT
# ============================================================
def rotate_backups():
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    for f in [TRAIN_PKL, VAL_PKL, TMP_PKL]:
        if os.path.exists(f):
            new_name = f.replace(".pkl", f"_{timestamp}.bak").replace(".tmp", f"_{timestamp}.bak")
            os.rename(f, new_name)
            log_message(f"Old file backed up → {new_name}")

    # Keep only latest backup
    for file in sorted(os.listdir(BASE_DIR)):
        if file.endswith(".bak"):
            file_path = os.path.join(BASE_DIR, file)
            age_hours = (time.time() - os.path.getmtime(file_path)) / 3600
            if age_hours > 24:  # older than 1 day
                os.remove(file_path)
                log_message(f"Deleted old backup → {file}")

# ============================================================
# MAIN FUNCTION
# ============================================================
def build_token_cache():
    log_message("===== Starting MKD Token Cache Builder v6 =====")
    log_message(f"Tokenizer: {TOKENIZER_PATH}")
    log_message(f"Dataset: {TEXT_PATH}")

    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
    eot = tokenizer.token_to_id("<|endoftext|>")
    log_message(f"Tokenizer loaded successfully. EOT id = {eot}")

    # Resume progress
    start_line = 0
    if os.path.exists(TMP_PKL):
        with open(TMP_PKL, "rb") as f:
            saved = pickle.load(f)
            start_line = saved.get("line", 0)
        log_message(f"Resuming from line {start_line:,}")
    else:
        log_message("No checkpoint found. Starting fresh.")
        rotate_backups()

    # Count dataset
    with open(TEXT_PATH, "r", encoding="utf-8") as f:
        total_lines = sum(1 for _ in f)
    split_index = int(total_lines * SPLIT_RATIO)
    log_message(f"Total lines: {total_lines:,} | Train split @ {split_index:,}")

    train_file = open(TRAIN_PKL, "ab")
    val_file = open(VAL_PKL, "ab")

    start_time = time.time()
    total_train_tokens, total_val_tokens = 0, 0
    batch_lines = []

    with open(TEXT_PATH, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i < start_line:
                continue

            line = line.strip()
            if not line:
                continue

            batch_lines.append(line)

            if len(batch_lines) >= BATCH_LINES:
                t_tok, v_tok = process_batch(tokenizer, batch_lines, eot, i, total_lines,
                                             split_index, train_file, val_file, start_time)
                total_train_tokens += t_tok
                total_val_tokens += v_tok
                batch_lines = []
                gc.collect()

                # token limit reached
                if total_train_tokens >= TOKEN_LIMIT_TRAIN and total_val_tokens >= TOKEN_LIMIT_VAL:
                    log_message("Target token limit reached — stopping safely.")
                    break

            if (i + 1) % SAVE_EVERY == 0:
                with open(TMP_PKL, "wb") as tmp:
                    pickle.dump({"line": i}, tmp)
                log_message(f"Checkpoint saved at line {i:,} | {system_status()}")

    if batch_lines:
        t_tok, v_tok = process_batch(tokenizer, batch_lines, eot, total_lines, total_lines,
                                     split_index, train_file, val_file, start_time)
        total_train_tokens += t_tok
        total_val_tokens += v_tok

    train_file.close()
    val_file.close()

    # Cleanup + new tmp
    if os.path.exists(TMP_PKL):
        os.remove(TMP_PKL)
    with open(TMP_PKL, "wb") as tmp:
        pickle.dump({"line": i}, tmp)
    log_message(f"New progress checkpoint saved at line {total_lines:,}")

    # Summary
    elapsed_hours = (time.time() - start_time) / 3600
    log_message("========== SUMMARY ==========")
    log_message(f"Train tokens generated: {total_train_tokens:,}")
    log_message(f"Val tokens generated:   {total_val_tokens:,}")
    log_message(f"Train cache → {TRAIN_PKL}")
    log_message(f"Val cache   → {VAL_PKL}")
    log_message(f"Elapsed time: {elapsed_hours:.2f} hours")
    log_message("=============================")

# ============================================================
# PROCESS BATCH
# ============================================================
def process_batch(tokenizer, lines, eot, current_line, total_lines,
                  split_index, train_file, val_file, start_time):
    encoded = tokenizer.encode_batch(lines)
    train_tokens, val_tokens = [], []
    for j, enc in enumerate(encoded):
        line_idx = current_line - len(encoded) + j + 1
        if line_idx < split_index:
            train_tokens.append(enc.ids + [eot])
        else:
            val_tokens.append(enc.ids + [eot])

    train_tok_count = sum(len(x) for x in train_tokens)
    val_tok_count = sum(len(x) for x in val_tokens)

    for seq in train_tokens:
        pickle.dump(seq, train_file)
    for seq in val_tokens:
        pickle.dump(seq, val_file)

    elapsed = (time.time() - start_time) / 60
    percent = (current_line / total_lines) * 100
    remaining = (elapsed / (percent + 1e-5)) * (100 - percent)
    msg = (f"Progress {percent:5.2f}% | Line {current_line:,}/{total_lines:,} "
           f"| TrainTok +{train_tok_count:,} | ValTok +{val_tok_count:,} "
           f"| Elapsed {elapsed:.1f}m | ETA {remaining/60:.1f}h | {system_status()}")
    log_message(msg)

    return train_tok_count, val_tok_count

# ============================================================
# AUTO RESTART LOOP
# ============================================================
if __name__ == "__main__":
    while True:
        try:
            build_token_cache()
            break
        except KeyboardInterrupt:
            log_message("Stopped manually. Progress saved.")
            sys.exit(0)
        except Exception as e:
            log_message(f"Error: {e}")
            log_message(f"Restarting automatically in {RESTART_DELAY}s...")
            for t in range(RESTART_DELAY, 0, -1):
                print(f"\rRestart in {t:02d}s | {system_status()}", end="", flush=True)
                time.sleep(1)
            print("\nRestarting now...\n")