SaylorTwift HF Staff

Resilient worker (transient CaptureContractError -> unscorable, systematic still fatal), resume-from-repo, eval worker base-cells download fix

e7168a9 verified 14 days ago

Raw History Blame Contribute Delete

8.16 kB

#!/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 ==="

	#!/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 ==="