#!/usr/bin/env bash
# Make the FP8 copy of Clef-Flash and compare it with BF16, end to end, on ML-Lab.
#
#   1. Quantize on CPU in a throwaway container (quantize_fp8.py). This step
#      is skipped if the FP8 directory already has a MANIFEST.json.
#   2. Start the BF16 model server, send the parity requests, save the answers
#      and the server log, and stop the server.
#   3. The same for the FP8 copy.
#   4. Compare the answers and write one report: results/<time>/report.md.
#
# Run it from anywhere:   ~/clef-fp8/run_all.sh
# Change a setting with an environment variable, for example:
#   PORT=8090 EXTRA_REQUESTS=~/real.jsonl ~/clef-fp8/run_all.sh
#
# The exit code is 0 for PASS and 1 for FAIL or for an error.
set -euo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
MODEL_DIR="${MODEL_DIR:-$HOME/clef-flash}"                     # Cloudflare's files (BF16)
FP8_DIR="${FP8_DIR:-$HOME/clef-flash-fp8}"                     # the FP8 copy goes here
# The model-server image for both servers. Set SERVER_IMAGE="" to run the code
# in CODE_DIR on the vLLM base image instead.
SERVER_IMAGE="${SERVER_IMAGE-docker-registry.wikimedia.org/wikimedia/machinelearning-liftwing-inference-services-decision-models-clef:2026-10-14-101821-publish}"
CODE_DIR="${CODE_DIR:-$HOME/inference-services}"               # used only when SERVER_IMAGE=""
PORT="${PORT:-8080}"                                           # gRPC uses PORT + 1
GPU="${GPU:-0}"                                                # GPU index, as rocm-smi numbers it
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.5}"        # share of the GPU for vLLM
EXTRA_REQUESTS="${EXTRA_REQUESTS:-}"                           # optional JSONL file of extra requests
BF16_URL="${BF16_URL:-}"            # use a BF16 server that is already running, e.g. http://localhost:8080/v1/models/clef-flash:predict
BF16_CONTAINER="${BF16_CONTAINER:-}"  # its container name, to read vLLM's memory numbers from its log (optional)
PROXY="${PROXY:-http://webproxy:8080}"
VLLM_BASE="${VLLM_BASE:-docker-registry.wikimedia.org/ml/amd-vllm022:gfx90agfx942rocm7.2.0pytorch2.10.0flash-attn2.8.3aiter0.1.13vllm0.22.1-3}"
RESULTS="${RESULTS:-$HERE/results/$(date -u +%Y%m%dT%H%M%SZ)}"
URL="http://localhost:$PORT/v1/models/clef-flash:predict"

fail() { echo "ERROR: $*" >&2; exit 1; }

# Stop early if the GPU does not have the memory that vLLM asks for. ML-Lab's
# GPU is shared. Skipped if rocm-smi is missing or prints another format.
check_gpu() {
  local used need
  used=$(rocm-smi --showmemuse 2>/dev/null | sed -n "s/^GPU\[$GPU\].*Allocated (VRAM%): *\([0-9][0-9]*\).*/\1/p" | head -1 || true)
  [ -n "$used" ] || return 0
  need=$(python3 -c "print(round($GPU_MEMORY_UTILIZATION * 100))")
  echo "   GPU $GPU: ${used}% of memory in use, vLLM needs ${need}% free"
  [ $((used + need)) -le 100 ] || fail "GPU $GPU does not have ${need}% free memory. See what uses it: rocm-smi --showpidgpus; docker ps. Or use GPU=<index> or a lower GPU_MEMORY_UTILIZATION."
}
cleanup() { docker rm -f clef-parity-bf16 clef-parity-fp8 >/dev/null 2>&1 || true; }
trap cleanup EXIT

# --- Checks before anything slow -------------------------------------------
[ -f "$MODEL_DIR/joint_head.safetensors" ] || fail "MODEL_DIR=$MODEL_DIR has no joint_head.safetensors."
[ -n "$SERVER_IMAGE" ] || grep -q collapse_image_pads "$CODE_DIR/src/models/decision_models/clef/model.py" 2>/dev/null \
  || fail "CODE_DIR=$CODE_DIR does not have the vLLM clef server."
python3 -c 'import sys; sys.exit(sys.version_info < (3, 9))' || fail "parity.py needs Python 3.9 or newer."
for p in "$PORT" "$((PORT + 1))"; do
  ! ss -ltn 2>/dev/null | grep -E ":${p}[[:space:]]" >/dev/null || fail "Port $p is in use. Run with PORT=<another port>."
done
[ -z "$EXTRA_REQUESTS" ] || [ -f "$EXTRA_REQUESTS" ] || fail "EXTRA_REQUESTS=$EXTRA_REQUESTS does not exist."
mkdir -p "$RESULTS"
echo "Results: $RESULTS"

# --- 1. Quantize -------------------------------------------------------------
PARTIAL="$FP8_DIR.partial"
if [ -f "$FP8_DIR/MANIFEST.json" ]; then
  echo "== 1/4 The FP8 copy exists ($FP8_DIR). Quantization skipped."
else
  [ ! -e "$FP8_DIR" ] || [ -z "$(ls -A "$FP8_DIR")" ] \
    || fail "$FP8_DIR has files but no MANIFEST.json, so it is not a finished FP8 copy. Remove it with: rm -rf $FP8_DIR"
  echo "== 1/4 Quantizing on CPU. This takes about 15 minutes and 35 GB of RAM."
  # Write to a separate directory, and move it into place only on success. A
  # failed run then leaves nothing half-written in FP8_DIR.
  rm -rf "$PARTIAL" && mkdir -p "$PARTIAL" && chmod 777 "$PARTIAL"
  # The constraint keeps the image's torch and transformers: pip fails rather
  # than replace them. pip also installs compressed-tensors 0.18.0, while the
  # vLLM in this image pins 0.15.0.1. This container does not run vLLM, so
  # --no-warn-conflicts hides that warning. No GPU devices, so llm-compressor
  # works on the CPU.
  docker run --rm --network host -e http_proxy="$PROXY" -e https_proxy="$PROXY" \
    -v "$MODEL_DIR":/src:ro -v "$PARTIAL":/dst -v "$HERE":/scripts:ro "$VLLM_BASE" bash -c '
      python3 -c "import torch, transformers; print(f\"torch=={torch.__version__}\ntransformers=={transformers.__version__}\")" > /tmp/keep.txt &&
      pip install -q --no-warn-conflicts -c /tmp/keep.txt llmcompressor==0.13.0 &&
      python3 /scripts/quantize_fp8.py --src /src --dst /dst' 2>&1 | tee "$RESULTS/quantize.log" \
    || fail "Quantization failed. Log: $RESULTS/quantize.log"
  # The container wrote the files as its own user, and some only readable by
  # that user. Give them to you, and make them readable for the model servers.
  docker run --rm --user 0 -v "$PARTIAL":/dst "$VLLM_BASE" \
    bash -c "chown -R $(id -u):$(id -g) /dst && chmod -R a+rX /dst" \
    || fail "Could not change the owner of the files in $PARTIAL."
  rmdir "$FP8_DIR" 2>/dev/null || true
  mv "$PARTIAL" "$FP8_DIR"
fi
cp "$FP8_DIR/MANIFEST.json" "$RESULTS/"

# --- 2 and 3. Serve each model and ask the same questions ----------------------
ask() {  # usage: ask LABEL MODEL_DIRECTORY STEP
  local label=$1 dir=$2 step=$3 name="clef-parity-$1" waited=0
  echo "== $step/4 Starting the $label server (vLLM start and warmup take a few minutes)"
  docker rm -f "$name" >/dev/null 2>&1 || true
  check_gpu
  local args=(--name "$name" --network host
    --device=/dev/kfd --device=/dev/dri --group-add video --group-add 105 --shm-size 8g
    -e HIP_VISIBLE_DEVICES="$GPU" -e MODEL_NAME=clef-flash -e MODEL_PATH=/mnt/models -e HF_HUB_OFFLINE=1
    -e GPU_MEMORY_UTILIZATION="$GPU_MEMORY_UTILIZATION" -v "$dir":/mnt/models:ro)
  if [ -n "$SERVER_IMAGE" ]; then
    # The image's entrypoint runs model.py and passes these arguments to it.
    docker run -d "${args[@]}" "$SERVER_IMAGE" --http_port "$PORT" --grpc_port "$((PORT + 1))" >/dev/null \
      || fail "Docker could not start the $label container."
  else
    docker run -d "${args[@]}" -e http_proxy="$PROXY" -e https_proxy="$PROXY" -e no_proxy=localhost,127.0.0.1 \
      -v "$CODE_DIR":/srv/app "$VLLM_BASE" \
      bash -c "cd /srv/app/src/models/decision_models/clef && pip install -q -r requirements.txt && \
               python3 model.py --http_port $PORT --grpc_port $((PORT + 1))" >/dev/null \
      || fail "Docker could not start the $label container."
  fi
  # grep -c, not grep -q: grep -q stops at the first match, `docker logs` then
  # fails with SIGPIPE, and under pipefail the check never passes.
  until docker logs "$name" 2>&1 | grep -c "Warmup complete" >/dev/null; do
    if [ "$(docker inspect -f '{{.State.Running}}' "$name")" != true ] || [ "$waited" -ge 1800 ]; then
      docker logs "$name" >"$RESULTS/server-$label.log" 2>&1
      echo "--- Error lines in the $label server log, first one first:"
      grep -E "Error|ERROR" "$RESULTS/server-$label.log" | cut -c1-240 | head -15 || true
      fail "The $label server did not start. Full log: $RESULTS/server-$label.log"
    fi
    sleep 10
    waited=$((waited + 10))
    [ $((waited % 60)) -ne 0 ] || echo "   still starting ($((waited / 60)) min): $(docker logs "$name" 2>&1 | tail -1 | cut -c1-100)"
  done
  echo "   Sending the parity requests to the $label server"
  https_proxy="$PROXY" python3 "$HERE/parity.py" run --label "$label" --out "$RESULTS/$label.json" \
    --url "$URL" --image-dir "$HERE/parity-images" ${EXTRA_REQUESTS:+--requests "$EXTRA_REQUESTS"} | tail -1 \
    || fail "Sending the parity requests to the $label server failed (see the error above)."
  docker logs "$name" >"$RESULTS/server-$label.log" 2>&1
  docker rm -f "$name" >/dev/null
}
if [ -n "$BF16_URL" ]; then
  echo "== 2/4 Using the BF16 server that is already running: $BF16_URL"
  curl -sf --noproxy '*' "${BF16_URL%:predict}" >/dev/null || fail "No ready model at ${BF16_URL%:predict}."
  https_proxy="$PROXY" python3 "$HERE/parity.py" run --label bf16 --out "$RESULTS/bf16.json" \
    --url "$BF16_URL" --image-dir "$HERE/parity-images" ${EXTRA_REQUESTS:+--requests "$EXTRA_REQUESTS"} | tail -1 \
    || fail "Sending the parity requests to the BF16 server failed (see the error above)."
  if [ -n "$BF16_CONTAINER" ]; then
    docker logs "$BF16_CONTAINER" >"$RESULTS/server-bf16.log" 2>&1 || echo "   Could not read the log of $BF16_CONTAINER."
  fi
else
  ask bf16 "$MODEL_DIR" 2
fi
ask fp8 "$FP8_DIR" 3

# --- 4. Report -----------------------------------------------------------------
echo "== 4/4 Writing the report"
python3 "$HERE/parity.py" report "$RESULTS" \
  --note "Servers: ${SERVER_IMAGE:-the code in $CODE_DIR on $VLLM_BASE}, on $(hostname), GPU $GPU, GPU_MEMORY_UTILIZATION=$GPU_MEMORY_UTILIZATION, one request at a time." \
  ${BF16_URL:+--note "The BF16 server was already running ($BF16_URL), so its settings are whatever it was started with."}