#!/usr/bin/env bash
# Per-checkpoint eval worker (1 GPU). Polls ADAPTER_REPO for new checkpoint-N
# dirs and evaluates each on a fixed N_TASKS subset of the held-out test split,
# writing eval/<checkpoint-N>_results.json + a rolling eval/summary.json.
#
# Expected env:
#   HF_TOKEN       (secret)  hub access + sandbox Jobs
#   ADAPTER_REPO   (env)     repo the training job pushes checkpoints to
#   N_TASKS        (env)     test tasks per eval (default 100; same subset every time)
set -euo pipefail
cd "$(dirname "$0")"

echo "=== [setup] ==="
apt-get update -qq && apt-get install -y -qq --no-install-recommends git ca-certificates
python3 -m pip install --quiet \
    'openenv[harbor]' \
    'harbor[hf-sandbox]>=0.24.0' \
    math-verify tomli-w \
    aiohttp
git clone --depth 1 https://github.com/huggingface/OpenEnv.git /opt/OpenEnv
SITE=$(python3 -c "import site; print(site.getsitepackages()[0])")
echo "/opt/OpenEnv/envs" > "$SITE/openenv_examples.pth"

echo "=== [patch] opencode seam: session id as x-api-key (Jobs proxy strips Authorization) ==="
python3 - <<'EOF'
import openenv.harbor.seams as seams

path = seams.__file__
src = open(path).read()
old = '"apiKey": session,\n                        "timeout": 600_000,'
new = ('"apiKey": session,\n'
       '                        "headers": {"x-api-key": session},\n'
       '                        "timeout": 600_000,')
assert src.count(old) == 1, f"opencode seam not found in {path}"
open(path, "w").write(src.replace(old, new))
print("patched", path)
EOF

echo "=== [setup] pinned test tasks ==="
python3 prepare_tasks.py --output /app/prepared --splits test

echo "=== [start] vLLM (base; LoRAs hot-loaded per checkpoint) ==="
CUDA_VISIBLE_DEVICES=0 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \
vllm serve "${BASE_MODEL:-LiquidAI/LFM2.5-2.6B}" \
    --port 8001 \
    --max-model-len 32768 \
    --return-tokens-as-token-ids \
    --logprobs-mode processed_logprobs \
    --enable-auto-tool-choice \
    --tool-call-parser lfm2 \
    --enable-lora --max-lora-rank 32 --max-loras 8 &
VLLM_PID=$!

echo "=== [start] env server (test split; tunnel-free via Jobs proxy) ==="
export SPACE_HOST="${JOB_ID}--8000.hf.jobs"
PREPARED=/app/prepared OPENENV_DATASETS=/app/prepared/datasets/test bash serve_env.sh &
ENV_PID=$!

python3 - <<'EOF'
import time, urllib.request
for name, url in [("vllm", "http://127.0.0.1:8001/v1/models"),
                  ("env",  "http://127.0.0.1:8000/health")]:
    for i in range(360):
        try:
            urllib.request.urlopen(url, timeout=5)
            print(name, "ready")
            break
        except Exception:
            time.sleep(5)
    else:
        raise SystemExit(f"{name} server never became ready")
EOF

echo "=== [watch] evaluating every new checkpoint in ${ADAPTER_REPO} ==="
python3 eval_watch.py

kill $VLLM_PID $ENV_PID 2>/dev/null || true
echo "=== done ==="