#!/usr/bin/env bash
# SUM Parts - unattended training with automatic resume
#
# The WSL VM has died twice under load on this box, so an overnight run needs to
# survive the process disappearing. main.py checkpoints every epoch and supports
#
#     mode=resume --pretrained_path <run>/checkpoint/<run>_ckpt_latest.pth
#
# When pretrained_path sits inside a checkpoint/ directory, resume_exp_directory
# reuses the SAME run folder and resume_checkpoint restores epoch, optimizer and
# scheduler -- so a restart continues rather than starting over.
#
# This wrapper runs the training, and if it exits non-zero it waits, finds the
# newest checkpoint, and resumes. Up to MAX_RETRIES times. It stops retrying if
# the run reached the final epoch, or if two consecutive attempts fail without
# a new checkpoint being written (that means it is failing before it can train,
# so retrying is pointless).
#
# Usage:
#   bash train_watchdog.sh [cfg]
#   MAX_RETRIES=10 EPOCHS=100 bash train_watchdog.sh pointnet
#
# For an unattended launch use launch_overnight.sh, which detaches this from
# the terminal.
set -uo pipefail

CONDA_ROOT="$HOME/miniconda3"
ENV_NAME="sumparts"
SEG="$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle/examples/segmentation"
DATA="$HOME/sum-parts/data/face_labeling/texsp_pcl"

CFG="${1:-pointnet}"
EPOCHS="${EPOCHS:-100}"
VAL_FREQ="${VAL_FREQ:-5}"
MAX_RETRIES="${MAX_RETRIES:-8}"
RETRY_WAIT="${RETRY_WAIT:-60}"

RUNROOT="$HOME/sum-parts/runs"
STAMP="${STAMP:-$(date +%Y%m%d-%H%M%S)}"
WORKDIR="$RUNROOT/${CFG}_${STAMP}"
LOG="$WORKDIR/watchdog.log"
TRAINLOG="$WORKDIR/train.log"
STATUS="$WORKDIR/status.txt"

if [ -n "${CFG_VOXEL_MAX:-}" ]; then
    VOXEL_MAX="$CFG_VOXEL_MAX"
else
    case "$CFG" in
        pointnext-xl)   VOXEL_MAX=32000 ;;
        pointvector-xl) VOXEL_MAX=24000 ;;
        *)              VOXEL_MAX=64000 ;;
    esac
fi

mkdir -p "$WORKDIR"

say() { echo "[$(date '+%F %T')] $*" | tee -a "$LOG"; }

write_status() {
    {
        echo "cfg        : $CFG"
        echo "workdir    : $WORKDIR"
        echo "state      : $1"
        echo "attempt    : ${attempt:-0}/$MAX_RETRIES"
        echo "epochs     : $EPOCHS  (val_freq $VAL_FREQ)"
        echo "voxel_max  : $VOXEL_MAX"
        echo "updated    : $(date '+%F %T')"
        echo "last epoch : $(last_epoch)"
        echo "train log  : $TRAINLOG"
    } > "$STATUS"
}

# newest *_ckpt_latest.pth under the segmentation log tree
find_ckpt() {
    find "$SEG/log/sumv2_triangle" -name '*_ckpt_latest.pth' -newermt "@$START_EPOCH_TS" \
        -printf '%T@ %p\n' 2>/dev/null | sort -rn | head -1 | cut -d' ' -f2-
}

last_epoch() {
    grep -oE 'Epoch [0-9]+ LR' "$TRAINLOG" 2>/dev/null | tail -1 | grep -oE '[0-9]+' || echo "-"
}

source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"
export WANDB_MODE=disabled WANDB_SILENT=true CUDA_HOME="$CONDA_PREFIX"

# Why expandable_segments: cfgs/sumv2_triangle/default.yaml validates with
#     val: { voxel_max: null }
# i.e. the whole tile (~470k points) in one go, while training is capped at
# voxel_max. Those oversized transient allocations fragment the caching
# allocator, its pool grows past 12 GB, and on WSL2 the driver quietly spills
# the excess into host RAM instead of raising OOM.
#
# Observed on the first run: epochs 1-10 took ~103 s each, then from epoch 11
# (right after the epoch-10 validation) every epoch took ~1100 s -- a 10x
# slowdown with the GPU at 100% util, full clocks, no thermal throttle, and
# 11.7 GB dedicated + 19.0 GB shared. Killing the process dropped the card back
# to 1.2 GB, so it was the trainer's own pool, not other apps.
#
# NOTE: expandable_segments is NOT available here -- it landed in torch 2.1 and
# this env is on 2.0.1, where it aborts at startup with
#     RuntimeError: Unrecognized CachingAllocator option: expandable_segments
# garbage_collection_threshold and max_split_size_mb do exist in 2.0.
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-garbage_collection_threshold:0.7,max_split_size_mb:128}"

# Validation runs with voxel_max: null, i.e. a whole ~470k-point tile in one
# forward pass, which peaks far above the capped training step. Bounding it to
# the training value keeps the allocator pool inside VRAM.
#
# This is a real deviation, stated plainly: val_miou is then computed on
# subsampled tiles, so it is only a model-selection signal. It does not affect
# the reported test numbers -- test() slides over the full tile regardless.
VAL_VOXEL_MAX="${VAL_VOXEL_MAX:-64000}"

if [ ! -d "$DATA/train" ]; then
    say "FATAL: $DATA/train missing. run download_data.sh + prepare_full_split.sh"
    write_status "failed-no-data"
    exit 1
fi

# RESUME_CKPT lets a relaunch continue an earlier run instead of starting over.
# Needed because find_ckpt only considers checkpoints written after this
# watchdog started, so a fresh watchdog would otherwise ignore existing ones.
RESUME_CKPT="${RESUME_CKPT:-}"
if [ -n "$RESUME_CKPT" ] && [ ! -f "$RESUME_CKPT" ]; then
    say "FATAL: RESUME_CKPT does not exist: $RESUME_CKPT"
    write_status "failed-bad-resume-ckpt"
    exit 1
fi

START_EPOCH_TS=$(date +%s)
cd "$SEG"

say "cfg=$CFG epochs=$EPOCHS val_freq=$VAL_FREQ voxel_max=$VOXEL_MAX"
say "data=$DATA  train/val/test = $(find -L "$DATA/train" -name '*.ply' | wc -l)/$(find -L "$DATA/val" -name '*.ply' | wc -l)/$(find -L "$DATA/test" -name '*.ply' | wc -l)"
say "workdir=$WORKDIR"

attempt=0
prev_ckpt=""

while [ "$attempt" -le "$MAX_RETRIES" ]; do
    ckpt=$(find_ckpt)
    # on the very first attempt, an explicitly supplied checkpoint wins
    if [ "$attempt" -eq 0 ] && [ -n "$RESUME_CKPT" ]; then
        ckpt="$RESUME_CKPT"
    fi

    if [ -z "$ckpt" ]; then
        say "attempt $attempt: fresh start"
        write_status "running (fresh)"
        python -u main.py \
            --cfg "../../cfgs/sumv2_triangle/${CFG}.yaml" \
            mode=train \
            dataset.common.data_root="$DATA" \
            dataset.train.voxel_max="$VOXEL_MAX" \
            dataset.val.voxel_max="$VAL_VOXEL_MAX" \
            epochs="$EPOCHS" \
            val_freq="$VAL_FREQ" \
            wandb.use_wandb=False \
            >> "$TRAINLOG" 2>&1
        rc=$?
    else
        say "attempt $attempt: resuming from $(basename "$ckpt")"
        write_status "running (resumed)"
        python -u main.py \
            --cfg "../../cfgs/sumv2_triangle/${CFG}.yaml" \
            mode=resume \
            --pretrained_path "$ckpt" \
            dataset.common.data_root="$DATA" \
            dataset.train.voxel_max="$VOXEL_MAX" \
            dataset.val.voxel_max="$VAL_VOXEL_MAX" \
            epochs="$EPOCHS" \
            val_freq="$VAL_FREQ" \
            wandb.use_wandb=False \
            >> "$TRAINLOG" 2>&1
        rc=$?
    fi

    if [ "$rc" -eq 0 ]; then
        say "training finished cleanly (rc=0) at epoch $(last_epoch)"
        write_status "done"
        say "last 20 lines:"
        tail -20 "$TRAINLOG" | tee -a "$LOG"
        exit 0
    fi

    say "attempt $attempt exited rc=$rc at epoch $(last_epoch)"
    tail -15 "$TRAINLOG" | tee -a "$LOG"

    new_ckpt=$(find_ckpt)
    if [ "$attempt" -gt 0 ] && [ "$new_ckpt" = "$prev_ckpt" ]; then
        say "no new checkpoint since the last attempt -- failing before training, giving up"
        write_status "failed-no-progress"
        exit "$rc"
    fi
    prev_ckpt="$new_ckpt"

    attempt=$((attempt + 1))
    if [ "$attempt" -gt "$MAX_RETRIES" ]; then
        say "exhausted $MAX_RETRIES retries"
        write_status "failed-retries-exhausted"
        exit "$rc"
    fi

    write_status "waiting ${RETRY_WAIT}s before retry"
    say "waiting ${RETRY_WAIT}s, then retrying"
    sleep "$RETRY_WAIT"

    # if the VM itself restarted, the GPU may take a moment to be usable again
    for i in $(seq 1 10); do
        if nvidia-smi -L >/dev/null 2>&1; then break; fi
        say "  GPU not ready yet ($i/10)"
        sleep 15
    done
done