#!/usr/bin/env bash
# Single-job entrypoint for HF Jobs (2-GPU flavor, e.g. h200x2).
# Runs all three processes in one container:
#   GPU 1: vLLM policy server          (serve_policy.sh)
#   CPU:   OpenEnv Harbor env server   (serve_env.sh)
#   GPU 0: TRL AsyncGRPOTrainer        (train_lora.py)
# Sandboxes: Harbor's hf-sandbox backend — each rollout boots an HF Jobs
# cpu-basic machine under the same HF token. No Daytona/E2B account needed.
#
# Expected env (via `hf jobs run --secrets/--env`):
#   HF_TOKEN          (secret)  hub downloads, checkpoint pushes, sandbox Jobs
#   ADAPTER_REPO      (env)     Hub repo that receives adapter checkpoints
#   MAX_STEPS         (env)     20 for the smoke run, 1000 for the full run
#   SANDBOX           (env)     hf-sandbox (default) | daytona | e2b
set -euo pipefail
cd "$(dirname "$0")"

echo "=== [setup] git (vllm image has none) ==="
apt-get update -qq && apt-get install -y -qq --no-install-recommends git ca-certificates

# NOTE: openenv pulls fastmcp 3.x -> mcp 1.x, which conflicts with vllm 0.32's
# mcp>=2. pip keeps mcp 1.x and only warns; if vllm serve crashes on an mcp
# import, pin an older image tag (vllm>=0.22 suffices for TRL's weight transfer).
echo "=== [setup] packages ==="
# vllm + torch come from the base image (vllm/vllm-openai). Per TRL docs,
# transformers>=5.2.0 may need --no-deps next to vllm's pinned stack.
python3 -m pip install --quiet 'transformers>=5.2.0' --no-deps
# trl is pinned to 1.15.0: 1.16.0 made tree packing unconditional in
# AsyncGRPOTrainer and rejects hybrid models (LFM2.5 has conv layers). 1.15.0
# has HarnessRolloutWorker (#6947) but predates tree packing.
python3 -m pip install --quiet \
    'trl[peft]==1.15.0' \
    'openenv[harbor]' \
    'harbor[hf-sandbox]>=0.24.0' \
    'kernels==0.16.0' \
    math-verify tomli-w \
    aiohttp

echo "=== [setup] OpenEnv repo (harbor_env lives in envs/, not the wheel) ==="
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 "=== [setup] pinned tasks (1,000 train / 250 test, hash-checked) ==="
python3 prepare_tasks.py --output /app/prepared

echo "=== [check] imports ==="
python3 - <<'EOF'
import trl, openenv, harbor_env  # noqa: F401
from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker  # noqa: F401
print("trl", trl.__version__, "| openenv", openenv.__version__)
EOF

echo "=== [patch] opencode seam: also send the capture session as x-api-key ==="
# The HF Jobs proxy strips Authorization. The capture proxy (sessions.py) also
# accepts x-api-key / x-goog-api-key, and OpenCode's @ai-sdk/openai-compatible
# provider supports a custom headers option — so send the session id both ways.
python3 - <<'EOF'
import openenv.harbor.seams as seams

path = seams.__file__
src = open(path).read()
# Target only the _opencode seam: its apiKey/timeout pair at 24-space indent is
# unique (other harnesses share the strings at different indents).
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 "=== [start] vLLM policy server (GPU 1) ==="
CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 \
vllm serve "${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 \
    --weight-transfer-config '{"backend":"nccl"}' &
VLLM_PID=$!

echo "=== [start] env server (CPU) ==="
# Tunnel-free proxy exposure: with SPACE_HOST set, OpenEnv mounts the capture
# proxy at /capture on the env server itself, reached through the Jobs proxy at
# https://<JOB_ID>--8000.hf.jobs — no gradio tunnel to die mid-run
# (that failure stalled the first full run silently for 13h).
export SPACE_HOST="${JOB_ID}--8000.hf.jobs"
echo "capture proxy public base: https://${SPACE_HOST}/capture"
PREPARED=/app/prepared bash serve_env.sh &
ENV_PID=$!

echo "=== [start] watchdog (fail the job if no training step for 30 min) ==="
(
  while true; do
    sleep 300
    if [ -f /tmp/rl-heartbeat ]; then
      age=$(( $(date +%s) - $(stat -c %Y /tmp/rl-heartbeat) ))
      if [ "$age" -gt 1800 ]; then
        echo "WATCHDOG: no training step for ${age}s — rollout path is dead, failing the job"
        pkill -f train_lora.py || true
        sleep 30
        exit 1
      fi
    fi
  done
) &
WATCHDOG_PID=$!

echo "=== [wait] servers ==="
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):  # up to 30 min for model load
        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 "=== [resume] latest pushed checkpoint, if any ==="
# Optimizer shards are not uploaded, so resume restores the policy, scheduler,
# task position (rollout_state.json) and RNG of the sampler — with a fresh
# optimizer state. Close enough for RL.
RESUME_FROM=""
if [ -n "${ADAPTER_REPO:-}" ]; then
  RESUME_FROM=$(python3 - <<'EOF'
import os, re
from huggingface_hub import HfApi, snapshot_download

repo = os.environ.get("ADAPTER_REPO", "")
try:
    files = HfApi().list_repo_files(repo)
except Exception:
    files = []
ckpts = sorted(
    {f.split("/")[0] for f in files if f.startswith("checkpoint-") and f.endswith("adapter_config.json")},
    key=lambda s: int(s.rsplit("-", 1)[1]),
)
if ckpts:
    path = snapshot_download(repo, allow_patterns=[f"{ckpts[-1]}/*"])
    print(os.path.join(path, ckpts[-1]))
EOF
)
fi
echo "RESUME_FROM=${RESUME_FROM:-<fresh start>}"

echo "=== [train] GPU 0, MAX_STEPS=${MAX_STEPS:-1000} ==="
CUDA_VISIBLE_DEVICES=0 \
VLLM_URL=http://127.0.0.1:8001 \
ENV_URL=http://127.0.0.1:8000 \
TRAIN_SPLIT=/app/prepared/datasets/train \
RESUME_FROM="$RESUME_FROM" \
    python3 train_lora.py

if [ "${SKIP_EVAL:-0}" != "1" ] && [ -n "${ADAPTER_REPO:-}" ]; then
    echo "=== [eval] restarting vLLM with the trained adapter (GPU 1) ==="
    kill $VLLM_PID 2>/dev/null || true; sleep 10
    # Late-bind the eval code: a fix uploaded during training still lands here.
    mkdir -p /app-eval
    hf download "${CODE_REPO:-SaylorTwift/lfm-opencode-lora-rl-code}" \
        --local-dir /app-eval --include "eval_adapter.py" || cp /app/eval_adapter.py /app-eval/
    CUDA_VISIBLE_DEVICES=1 \
    vllm serve "${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 \
        --lora-modules "trained=${ADAPTER_REPO}" &
    VLLM_PID=$!
    python3 - <<'EOF'
import time, urllib.request
for i in range(360):
    try:
        urllib.request.urlopen("http://127.0.0.1:8001/v1/models", timeout=5)
        print("vllm ready")
        break
    except Exception:
        time.sleep(5)
else:
    raise SystemExit("vllm never became ready for eval")
EOF
    echo "=== [eval] 250 held-out tasks x {base, trained} ==="
    N_TASKS=250 OUT=/app/eval_results_final.json ADAPTER_NAME=trained \
        python3 /app-eval/eval_adapter.py
    python3 - <<'EOF'
import os
from huggingface_hub import HfApi
repo = os.environ.get("ADAPTER_REPO")
if repo:
    HfApi().upload_file(path_or_fileobj="/app/eval_results_final.json",
                        path_in_repo="eval/final_results.json",
                        repo_id=repo)
    print("uploaded to", repo)
EOF
fi

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