#!/usr/bin/env bash
set -euo pipefail

cd "$(dirname "$0")"

CONTAINER_NAME="llmc_codegpt_multilingual_gpu01"
IMAGE="nvidia/cuda:12.2.2-devel-ubuntu22.04"
RUN_ROOT="log_codegpt_multilingual_runs"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
RUN_DIR="${RUN_ROOT}/${TIMESTAMP}"
LATEST_LINK="${RUN_ROOT}/latest"
NCCL_FS_DIR="/tmp/llmc_nccl_codegpt_multi_${TIMESTAMP}"

EPOCHS=5
MICRO_BATCH_SIZE=32
SEQ_LEN=1024
NUM_PROCESSES=2
TOTAL_BATCH_SIZE=196608
LEARNING_RATE="0.0003"
FINAL_LR_FRAC="0.1"
WARMUP_STEPS=1000
WEIGHT_DECAY="0.1"
RECOMPUTE=1
VAL_MAX_STEPS=20

TRAIN_PATTERN="dev/data/code_messages/code_messages_train_*.bin"
VAL_PATTERN="dev/data/code_messages/code_messages_val_*.bin"
LOAD_MODEL="codegpt_multilingual.bin"

read -r TRAIN_TOKENS STEPS_PER_EPOCH MAX_STEPS <<<"$(python - <<'PY'
import glob
import struct

files = sorted(glob.glob("dev/data/code_messages/code_messages_train_*.bin"))
if not files:
    raise SystemExit("no training shards found")

tokens = 0
for path in files:
    with open(path, "rb") as f:
        header = struct.unpack("256i", f.read(1024))
    tokens += header[2]

total_batch_size = 196608
epochs = 5
steps_per_epoch = tokens // total_batch_size
print(tokens, steps_per_epoch, steps_per_epoch * epochs)
PY
)"

VAL_EVERY=$((STEPS_PER_EPOCH / 16))
CHECKPOINT_EVERY="$STEPS_PER_EPOCH"
SAMPLE_EVERY="$STEPS_PER_EPOCH"

mkdir -p "$RUN_DIR"
ln -sfn "$TIMESTAMP" "$LATEST_LINK"

docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true

{
  echo "host_launch_time=$(date -Is)"
  echo "container_name=$CONTAINER_NAME"
  echo "image=$IMAGE"
  echo "run_dir=$RUN_DIR"
  echo "train_tokens=$TRAIN_TOKENS"
  echo "epochs=$EPOCHS"
  echo "steps_per_epoch=$STEPS_PER_EPOCH"
  echo "max_steps=$MAX_STEPS"
  echo "checkpoint_every=$CHECKPOINT_EVERY"
  echo "sample_every=$SAMPLE_EVERY"
  echo "val_every=$VAL_EVERY"
  echo "val_max_steps=$VAL_MAX_STEPS"
  echo "total_batch_size=$TOTAL_BATCH_SIZE"
  echo "micro_batch_size=$MICRO_BATCH_SIZE"
  echo "seq_len=$SEQ_LEN"
  echo "learning_rate=$LEARNING_RATE"
  echo "final_lr_frac=$FINAL_LR_FRAC"
  echo "warmup_steps=$WARMUP_STEPS"
  echo "weight_decay=$WEIGHT_DECAY"
  echo "recompute=$RECOMPUTE"
  echo "load_model=$LOAD_MODEL"
  echo "train_pattern=$TRAIN_PATTERN"
  echo "val_pattern=$VAL_PATTERN"
  echo "git_commit=$(git rev-parse HEAD 2>/dev/null || true)"
} > "$RUN_DIR/host_launch_meta.txt"

docker run -d \
  --name "$CONTAINER_NAME" \
  --gpus '"device=0,1"' \
  --ipc=host \
  -u "$(id -u):$(id -g)" \
  -v /raid/data/weifeng/llm.c:/workspace \
  -w /workspace \
  -e RUN_DIR="$RUN_DIR" \
  -e NCCL_FS_DIR="$NCCL_FS_DIR" \
  -e TRAIN_TOKENS="$TRAIN_TOKENS" \
  -e EPOCHS="$EPOCHS" \
  -e STEPS_PER_EPOCH="$STEPS_PER_EPOCH" \
  -e MAX_STEPS="$MAX_STEPS" \
  -e CHECKPOINT_EVERY="$CHECKPOINT_EVERY" \
  -e SAMPLE_EVERY="$SAMPLE_EVERY" \
  -e VAL_EVERY="$VAL_EVERY" \
  -e VAL_MAX_STEPS="$VAL_MAX_STEPS" \
  -e MICRO_BATCH_SIZE="$MICRO_BATCH_SIZE" \
  -e SEQ_LEN="$SEQ_LEN" \
  -e TOTAL_BATCH_SIZE="$TOTAL_BATCH_SIZE" \
  -e LEARNING_RATE="$LEARNING_RATE" \
  -e FINAL_LR_FRAC="$FINAL_LR_FRAC" \
  -e WARMUP_STEPS="$WARMUP_STEPS" \
  -e WEIGHT_DECAY="$WEIGHT_DECAY" \
  -e RECOMPUTE="$RECOMPUTE" \
  -e LOAD_MODEL="$LOAD_MODEL" \
  -e TRAIN_PATTERN="$TRAIN_PATTERN" \
  -e VAL_PATTERN="$VAL_PATTERN" \
  "$IMAGE" \
  bash -lc '
    set -euo pipefail

    mkdir -p "$RUN_DIR" "$NCCL_FS_DIR"
    rm -f "$NCCL_FS_DIR"/*

    {
      echo "container_start_time=$(date -Is)"
      echo "hostname=$(hostname)"
      echo "run_dir=$RUN_DIR"
      echo "train_tokens=$TRAIN_TOKENS"
      echo "epochs=$EPOCHS"
      echo "steps_per_epoch=$STEPS_PER_EPOCH"
      echo "max_steps=$MAX_STEPS"
      echo "checkpoint_every=$CHECKPOINT_EVERY"
      echo "sample_every=$SAMPLE_EVERY"
      echo "val_every=$VAL_EVERY"
      echo "val_max_steps=$VAL_MAX_STEPS"
      echo "total_batch_size=$TOTAL_BATCH_SIZE"
      echo "micro_batch_size=$MICRO_BATCH_SIZE"
      echo "seq_len=$SEQ_LEN"
      echo "learning_rate=$LEARNING_RATE"
      echo "final_lr_frac=$FINAL_LR_FRAC"
      echo "warmup_steps=$WARMUP_STEPS"
      echo "weight_decay=$WEIGHT_DECAY"
      echo "recompute=$RECOMPUTE"
      echo "load_model=$LOAD_MODEL"
      echo "train_pattern=$TRAIN_PATTERN"
      echo "val_pattern=$VAL_PATTERN"
      echo
      nvcc --version
      echo
      nvidia-smi
    } > "$RUN_DIR/run_meta.txt" 2>&1

    nvidia-smi \
      --query-gpu=timestamp,index,name,utilization.gpu,memory.used,memory.total,power.draw,temperature.gpu \
      --format=csv \
      -l 30 > "$RUN_DIR/gpu_metrics.csv" 2>&1 &
    monitor_pid=$!
    cleanup() {
      kill "$monitor_pid" 2>/dev/null || true
    }
    trap cleanup EXIT

    rm -f train_gpt2cu
    make train_gpt2cu PRECISION=FP32 FORCE_NVCC_O=3 > "$RUN_DIR/build.log" 2>&1

    common_args=(
      -e "$LOAD_MODEL"
      -i "$TRAIN_PATTERN"
      -j "$VAL_PATTERN"
      -o "$RUN_DIR"
      -b "$MICRO_BATCH_SIZE"
      -t "$SEQ_LEN"
      -d "$TOTAL_BATCH_SIZE"
      -x "$MAX_STEPS"
      -v "$VAL_EVERY"
      -m "$VAL_MAX_STEPS"
      -s "$SAMPLE_EVERY"
      -r "$RECOMPUTE"
      -c "$WEIGHT_DECAY"
      -l "$LEARNING_RATE"
      -q "$FINAL_LR_FRAC"
      -u "$WARMUP_STEPS"
      -n "$CHECKPOINT_EVERY"
      -f 0
      -y 0
      -pi fs
      -pf "$NCCL_FS_DIR"
      -pn 2
      -pg 2
    )

    printf "%q " ./train_gpt2cu "${common_args[@]}" -pr 0 > "$RUN_DIR/command_rank0.txt"
    printf "\n" >> "$RUN_DIR/command_rank0.txt"
    printf "%q " ./train_gpt2cu "${common_args[@]}" -pr 1 > "$RUN_DIR/command_rank1.txt"
    printf "\n" >> "$RUN_DIR/command_rank1.txt"

    ./train_gpt2cu "${common_args[@]}" -pr 0 > "$RUN_DIR/rank0.log" 2>&1 &
    pid0=$!
    ./train_gpt2cu "${common_args[@]}" -pr 1 > "$RUN_DIR/rank1.log" 2>&1 &
    pid1=$!

    set +e
    wait -n "$pid0" "$pid1"
    first_status=$?
    if [[ "$first_status" -ne 0 ]]; then
      kill "$pid0" "$pid1" 2>/dev/null || true
      wait "$pid0" "$pid1" 2>/dev/null
      exit "$first_status"
    fi
    wait "$pid0"
    status0=$?
    wait "$pid1"
    status1=$?
    if [[ "$status0" -ne 0 || "$status1" -ne 0 ]]; then
      exit 1
    fi
  '

echo "Launched $CONTAINER_NAME"
echo "Run dir: $RUN_DIR"
echo "Latest link: $LATEST_LINK"