#!/usr/bin/env bash
# gpu_preload_check.sh — GPU residency guard (design DESIGN-a2a-agents.md sec 3b)
#
# Never-evict set on the 8GB GTX 1070: bge-m3 (embedder) + tei-reranker.
# Evicting either silently breaks Hindsight recall (the memory plugin's
# recall timeout just skips injection, no error surfaced) — the whole
# reason this guard exists.
#
# Usage: gpu_preload_check.sh <requested_mib> [gpu_index]
#   requested_mib — VRAM footprint (MiB) of the model/process about to load
#   gpu_index     — nvidia-smi GPU index (default 0)
#
# Exit 0  — safe to proceed, never-evict set stays resident with headroom.
# Exit 1  — reject: loading this would eat into or evict the never-evict set.
# Exit 2  — reject: never-evict set isn't even currently resident (abort,
#           something is already wrong — don't compound it by loading more).
#
# This is a guard for callers (workers/scripts) that are about to pull a
# model onto the shared GPU. It does NOT itself load or evict anything.

set -euo pipefail

REQUESTED_MIB="${1:?usage: gpu_preload_check.sh <requested_mib> [gpu_index]}"
GPU_INDEX="${2:-0}"

# tei-reranker measured footprint (2026-07-26, jina-reranker-v2-base-multilingual
# fp16 on CUDA torch): ~1690 MiB resident. bge-m3 measured ~882 MiB via ollama.
# Keep these as a documented floor, not just "whatever's currently resident" —
# a transient dip during another process's own load shouldn't false-negative us.
RERANKER_FLOOR_MIB=1690
BGE_M3_FLOOR_MIB=882
NEVER_EVICT_FLOOR_MIB=$((RERANKER_FLOOR_MIB + BGE_M3_FLOOR_MIB))

log() { echo "[gpu_preload_check] $*" >&2; }

# 1. Confirm the never-evict set is actually resident right now.
reranker_up=0
if curl -fsS -m 3 "http://localhost:8014/info" >/dev/null 2>&1; then
  reranker_up=1
fi

bge_m3_up=0
if docker exec ollama ollama ps 2>/dev/null | grep -q '^bge-m3'; then
  bge_m3_up=1
fi

if [[ "$reranker_up" -ne 1 || "$bge_m3_up" -ne 1 ]]; then
  log "REJECT: never-evict set not fully resident (tei-reranker up=$reranker_up, bge-m3 up=$bge_m3_up)."
  log "Something is already wrong — fix that before loading anything else onto the GPU."
  exit 2
fi

# 2. Check free VRAM and whether the requested load would eat into the
#    never-evict floor.
free_mib=$(nvidia-smi --id="$GPU_INDEX" --query-gpu=memory.free --format=csv,noheader,nounits | tr -d ' ')

if [[ -z "$free_mib" ]]; then
  log "REJECT: could not read nvidia-smi free memory for GPU $GPU_INDEX."
  exit 1
fi

remaining_after_load=$((free_mib - REQUESTED_MIB))

log "free=${free_mib}MiB requested=${REQUESTED_MIB}MiB never_evict_floor=${NEVER_EVICT_FLOOR_MIB}MiB remaining_after_load=${remaining_after_load}MiB"

if (( remaining_after_load < 0 )); then
  log "REJECT: requested load (${REQUESTED_MIB}MiB) exceeds current free VRAM (${free_mib}MiB)."
  log "The kernel driver would have to evict something to fit it — on this box that means"
  log "risking the never-evict set (bge-m3 + tei-reranker). Refusing."
  exit 1
fi

log "OK: load fits in free VRAM without necessitating eviction of the never-evict set."
exit 0