#!/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}"
# Pins are only a fallback: the HF cache's `refs/main` is the truth (a
# re-cut/re-download moves the snapshot, and a hardcoded rev then points at a
# directory that no longer exists — exactly what broke the 2026-10-21 serve).
pin_or_ref() { # <hub dir> <pinned rev>
    local d="$NAS/$1" ref
    ref=$(cat "$d/refs/main" 2>/dev/null || true)
    if [[ -n "$ref" && -d "$d/snapshots/$ref" ]]; then echo "$ref"; else echo "$2"; fi
}
REV="${TF1017_MODEL_REV:-$(pin_or_ref models--LibertAIDAI--DeepSeek-V4.1-Flash-1017-EXL3-2.9bpw afc7aa0aab98ee735c0f54e0a122019af1ac08f1)}"
EREV="${TF1017_ENGRAM_REV:-$(pin_or_ref models--deepseek-ai--DeepSeek-V4.1-Flash-1017 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"
[[ -d "$WBUNDLE" ]] || { echo "[tf1017] FATAL: no checkpoint snapshot at $WBUNDLE" >&2; exit 1; }
[[ -d "$EBUNDLE" ]] || { echo "[tf1017] FATAL: no Engram source snapshot at $EBUNDLE" >&2; exit 1; }
echo "[tf1017] checkpoint rev $REV, Engram rev $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
        echo "[tf1017] packing Engram tables for rank $RANK (~94 GiB, reads 189 GiB from the NAS)"
        if [[ "${TF1017_PACK_IN_IMAGE:-1}" = "1" ]] && docker image inspect "$IMG" >/dev/null 2>&1; then
            # Prefer the image: it has numpy + the exact python the packer was
            # written against. Host python is PEP-668 managed on dgx7/8
            # (`error: externally-managed-environment`) — no host installs.
            # Mount the whole repo dir (`snapshots/` + `blobs/`), not just the
            # snapshot: the snapshot files are symlinks into the hub's blobs,
            # so a snapshot-only mount makes model.safetensors.index.json
            # dangle inside the container.
            # Mount the whole hub: snapshot files are symlinks to the
            # per-model blobs, which are themselves symlinks to hub-level
            # blobs (huggingface_hub >= 0.34 layout) — anything less leaves
            # model.safetensors.index.json dangling in the container.
            docker run --rm --network none --entrypoint python3 \
                -v "$NAS":/hub:ro -v "$CKPT":/ckpt:ro -v "$PACKED":/out \
                -v "$HELP/pack_engram.py":/pack_engram.py:ro \
                "$IMG" /pack_engram.py --src "/hub/${EBUNDLE#"$NAS"/}" \
                --config /ckpt/config.json --out /out --rank "$RANK"
        else
            python3 -c 'import numpy' 2>/dev/null \
                || python3 -m pip install --user --quiet --break-system-packages numpy \
                || { echo "[tf1017] FATAL: numpy missing for the Engram packer and pip refused (no image on this node yet)" >&2; exit 1; }
            python3 "$HELP/pack_engram.py" \
                --src "$EBUNDLE" --config "$CKPT/config.json" \
                --out "$PACKED" --rank "$RANK"
        fi
    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"