#!/bin/bash
# Serve Agnes 3.0 Flash with a stock sglang image.
#
# Tested image:  lmsysorg/sglang:nightly-dev-20260908-20ca564b
#
#   docker run --gpus all --shm-size 64g -p 30001:30002 \
#       -v /path/to/agnes-3.0-flash:/model \
#       lmsysorg/sglang:nightly-dev-20260908-20ca564b \
#       bash /model/serve.sh [extra sglang arguments, e.g. --tp 2]
#
# The script overlays three files from sglang_patch/ onto the image's sglang
# package (a prebuilt variant when one matches the installed version, else
# apply_patch.py patches the installed files in place) and starts the server
# on port 30002 inside the container.
set -euo pipefail
D="$(cd "$(dirname "$0")" && pwd)"
PKG="$(python3 -c 'import sglang, os; print(os.path.dirname(sglang.__file__))')"
VER="$(python3 -c 'import sglang; print(getattr(sglang, "__version__", ""))' 2>/dev/null || true)"

case "$VER" in
  *20ca564b*) VARIANT="nightly-dev-20260908-20ca564b" ;;
  0.5.19*)    VARIANT="v0.5.19" ;;
  *)          VARIANT="" ;;
esac

if [ -n "$VARIANT" ] && [ -d "$D/sglang_patch/$VARIANT/sglang" ]; then
  echo "[serve.sh] sglang $VER -> prebuilt patch $VARIANT"
  cp -r "$D/sglang_patch/$VARIANT/sglang/." "$PKG/"
else
  echo "[serve.sh] sglang $VER -> patching installed package in place"
  python3 "$D/sglang_patch/apply_patch.py" "$PKG"
fi

export AGNES_MODEL_PATH="$D"
exec python3 -m sglang.launch_server --model-path "$D" --trust-remote-code --host 0.0.0.0 --port 8080 "$@"