#!/bin/bash
# Run one end-to-end pass and report only what the pass is judged on.
#
# A successful run is not just "some text came out": it is every data-parallel rank agreeing
# on the token ids, and the fallback list matching what we already know is unported. Both are
# easy to skim past in the raw log, which prints one block per rank, so they are extracted here.
#
#   bash scripts/h20_e2e_report.sh
#   EP=8 MAX_TOKENS=64 bash scripts/h20_e2e_report.sh
#
# Env in: EP (default 8), MAX_TOKENS (default 32), LOG (default e2e.log),
# DEEPINFER_HF_MODEL_DIR (required -- the weights).
#
# Two settings are defaulted here because single-node H20 always needs them on this image, and
# both are explained in H20-TRITON-PLAN.md:
#   EP_DISABLE_GIN=1          no /dev/infiniband in the container; SM90 single node uses
#                             scale-up, so DeepEP should not look for GIN.
#   DEEPINFER_ENFORCE_EAGER=1 the Engram local gather does a D2H inside CUDA graph capture.

set -uo pipefail

EP="${EP:-8}"
MAX_TOKENS="${MAX_TOKENS:-32}"
LOG="${LOG:-e2e.log}"

export DEEPINFER_OPS_BACKEND="${DEEPINFER_OPS_BACKEND:-triton}"
export EP_DISABLE_GIN="${EP_DISABLE_GIN:-1}"
export DEEPINFER_ENFORCE_EAGER="${DEEPINFER_ENFORCE_EAGER:-1}"

# cu13 images: Triton downloads its own cudart but does not add it to the search path, so any
# kernel load fails with "libcudart.so.12: cannot open shared object file" without this.
for d in /root/.triton/nvidia/cudart/*/lib; do
    [ -d "$d" ] && export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}:$d"
done

if [ -z "${DEEPINFER_HF_MODEL_DIR:-}" ]; then
    echo "DEEPINFER_HF_MODEL_DIR is not set; point it at the weights directory" >&2
    exit 2
fi

python scripts/h20_ask.py --ep "$EP" --max_tokens "$MAX_TOKENS" >"$LOG" 2>&1
echo "exit=$?"

echo
echo "=== token ids: distinct sequence and how many ranks produced it"
# `uniq -c` is the point: a sequence appearing once per rank means every rank agreed.
grep -o "=== token ids ([0-9]*): \[[^]]*\]" "$LOG" | sed 's/=== token ids ([0-9]*): //' | sort | uniq -c

echo
echo "=== continuation, one per distinct prompt"
grep -o "=== continuation: '.*'" "$LOG" | sort -u

echo
echo "=== kernels served by the default backend"
grep -o "\[ops\] kernel '[a-z_0-9]*' used from '[a-z]*'" "$LOG" | sed "s/.*kernel '//;s/' used from.*//" | sort -u
grep -o "backend '[a-z]*' lacks [0-9]* kernel(s)" "$LOG" | sort -u
grep -q "rejection sampling" "$LOG" && echo "fused_sample: rejection-sampling variant (MTP verify)"

echo
echo "=== hard failures"
# Tracebacks specifically: the stats logger mentions "AttributeError" without the run failing.
if grep -q "Traceback" "$LOG"; then
    echo "TRACEBACK"
    grep -A 5 "Traceback" "$LOG" | head -20
else
    echo "none"
fi