#!/usr/bin/env bash
# tensorfold-1017-prepare — host-side staging for the TensorFold 2x 1017 lane.
#
# Idempotent: every step is skipped when its output already exists. Runs on
# each launch node before the containers start (recipe `before_serve`), with
# NODE_RANK exported by the launcher.
#
# Everything here works WITHOUT the image: the launcher runs pre-serve hooks
# *before* it builds/syncs the image, so the helper files (Engram packer, chat
# template, CUDA prebuild inputs) are published to the NAS by the recipe image's
# build.sh. The one image-dependent step (CUDA prebuild) is skipped when the
# image is not on this node yet and simply runs on the next serve.
#
# Steps
#   1. node-local copy of the EXL3 2.9bpw checkpoint (EXL3 streams weights on
#      every decode step — reading them off the NAS costs most of the decode);
#   2. the chat template the TensorFold HTTP app requires at init (the EXL3
#      checkpoint carries none);
#   3. the per-rank Engram packs (layers 1 + 14, ~47 GiB each) built from the
#      official 1017 shards 47/48 — rank r owns hash heads [12r, 12r+12);
#   4. the CUDA extension prebuild (CPU-only, 18/18 builders) when the image is
#      already present.
#
# Overrides (env): TF1017_ROOT (default /home/kilork/ns/1017-tf),
# TF1017_NAS (default /mnt/nas/ai/huggingface/hub), TF1017_IMAGE,
# TF1017_SKIP_WEIGHTS=1, TF1017_SKIP_PACK=1, TF1017_SKIP_PREBUILD=1.
set -euo pipefail

ROOT="${TF1017_ROOT:-/home/kilork/ns/1017-tf}"
NAS="${TF1017_NAS:-/mnt/nas/ai/huggingface/hub}"
PUB="${TF1017_PUB:-/mnt/nas/ai/sm121/dsv41-tf-1017}"
IMG="${TF1017_IMAGE:-dsv41-tf-1017:v1}"
RANK="${NODE_RANK:-0}"
REV="${TF1017_MODEL_REV:-afc7aa0aab98ee735c0f54e0a122019af1ac08f1}"
EREV="${TF1017_ENGRAM_REV:-1bae0e1088993b9b82de08f59de5be5a34e42170}"

WBUNDLE="$NAS/models--LibertAIDAI--DeepSeek-V4.1-Flash-1017-EXL3-2.9bpw/snapshots/$REV"
EBUNDLE="$NAS/models--deepseek-ai--DeepSeek-V4.1-Flash-1017/snapshots/$EREV"
CKPT="$ROOT/models/checkpoint"
PACKED="$ROOT/engram-packed-1017"
CACHE="$ROOT/cache"
HELP="$ROOT/helpers"

mkdir -p "$CKPT" "$PACKED" "$CACHE" "$HELP" "$ROOT/state"

echo "[tf1017] rank $RANK — root $ROOT"

# 0. helper files (published to the NAS by the image's build.sh) -------------
for f in pack_engram.py prebuild_ext.py prebuild-offline.py pfdense-table.json \
         chat_template.jinja; do
    if [[ ! -s "$HELP/$f" ]]; then
        if [[ -s "$PUB/$f" ]]; then
            cp -f "$PUB/$f" "$HELP/$f"
        elif docker image inspect "$IMG" >/dev/null 2>&1; then
            docker run --rm --network none --entrypoint cat "$IMG" "/opt/sm121-helpers/$f" > "$HELP/$f"
        else
            echo "[tf1017] FATAL: $f not published at $PUB and image $IMG absent" >&2
            exit 1
        fi
    fi
done

# 1. weights -----------------------------------------------------------------
if [[ "${TF1017_SKIP_WEIGHTS:-0}" != "1" ]]; then
    have=$(find "$CKPT" -maxdepth 1 -name '*.safetensors' | wc -l)
    if [[ "$have" -lt 43 ]]; then
        echo "[tf1017] copying EXL3 checkpoint from the NAS ($have/43 shards present)"
        tar -chf - -C "$WBUNDLE" . | tar -xf - -C "$CKPT"
    fi
    have=$(find "$CKPT" -maxdepth 1 -name '*.safetensors' | wc -l)
    [[ "$have" -ge 43 ]] || { echo "[tf1017] FATAL: checkpoint incomplete ($have/43)" >&2; exit 1; }
fi

# 2. chat template -----------------------------------------------------------
[[ -s "$CKPT/chat_template.jinja" ]] || cp "$HELP/chat_template.jinja" "$CKPT/chat_template.jinja"

# 3. per-rank Engram packs (host python + numpy; no image needed) ------------
if [[ "${TF1017_SKIP_PACK:-0}" != "1" ]]; then
    if [[ ! -s "$PACKED/engram-l14-r${RANK}of2.bin" ]]; then
        python3 -c 'import numpy' 2>/dev/null || python3 -m pip install --user --quiet numpy
        echo "[tf1017] packing Engram tables for rank $RANK (~94 GiB, reads 189 GiB from the NAS)"
        python3 "$HELP/pack_engram.py" \
            --src "$EBUNDLE" --config "$CKPT/config.json" \
            --out "$PACKED" --rank "$RANK"
    fi
    [[ -s "$PACKED/engram-l1-r${RANK}of2.bin" && -s "$PACKED/engram-l14-r${RANK}of2.bin" ]] \
        || { echo "[tf1017] FATAL: Engram packs missing for rank $RANK" >&2; exit 1; }
fi

# 4. CUDA extension prebuild (needs the image; skipped on a first serve) -----
if [[ "${TF1017_SKIP_PREBUILD:-0}" != "1" && ! -s "$CACHE/pfdense-table.json" ]]; then
    if docker image inspect "$IMG" >/dev/null 2>&1; then
        echo "[tf1017] prebuilding CUDA extensions (CPU-only)"
        docker run --rm --network none --memory 8g --cpus 4 \
            -e CUDA_VISIBLE_DEVICES= -e TORCH_CUDA_ARCH_LIST=12.1 -e MAX_JOBS=4 \
            -e TORCH_EXTENSIONS_DIR=/cache/torch_extensions \
            -v "$CACHE":/cache \
            -v "$HELP/prebuild-offline.py":/offline.py:ro \
            -v "$HELP/prebuild_ext.py":/prebuild.py:ro \
            -v "$HELP/pfdense-table.json":/table.json:ro \
            --entrypoint bash "$IMG" \
            -c 'python /offline.py && cp /table.json /cache/pfdense-table.json'
    else
        echo "[tf1017] image $IMG not on this node yet — extensions will compile at first start"
    fi
fi

echo "[tf1017] rank $RANK ready"