onf/scripts/run_cache_policy_actions.sh

#!/usr/bin/env bash
# Cache the frozen StableVLA policy's per-frame actions over one suite's demo corpus.
#
#   bash scripts/run_cache_policy_actions.sh <suite> [PORT] [extra args for the python script]
#
# Starts a policy server on PORT if nothing is listening there yet, runs scripts/cache_policy_actions.py
# against it, and tears the server back down if (and only if) this script started it -- a killed run
# that leaves one behind holds ~7.5 GiB of GPU memory until someone notices.
#
# ONF_GPU picks the device (default 0). OMP_NUM_THREADS=4 is not optional: leaving it unset once
# cost a 30s/it regression through OpenMP oversubscription.
#
# The whole body re-executes itself under setsid on first entry, so a dropped ssh session cannot take
# a multi-hour run with it. Everything after that point goes to logs/_drivers/, not the terminal.
set -uo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SUITE="${1:?usage: run_cache_policy_actions.sh <object|spatial|goal|long> [PORT] [extra args]}"
PORT="${2:-14190}"
shift $(( $# > 2 ? 2 : $# ))
LOG="$REPO/logs/_drivers/cache_policy_actions_${SUITE}.log"

if [ "${ONF_SETSID_DONE:-0}" != "1" ]; then
    mkdir -p "$(dirname "$LOG")"
    echo "[$(date -u +%FT%TZ)] detaching; log: $LOG"
    ONF_SETSID_DONE=1 ONF_GPU="${ONF_GPU:-0}" setsid "$BASH" "${BASH_SOURCE[0]}" "$SUITE" "$PORT" "$@" \
        >>"$LOG" 2>&1 </dev/null &
    exit 0
fi

PY="${ONF_PY:-/home/quang/miniconda3/envs/stablevla/bin/python}"
CONDA_PREFIX_DIR="$(dirname "$(dirname "$PY")")"
NVML_PRELOAD="${ONF_NVML_PRELOAD:-/srv/data/HAMLET-Isaac-GR00T/.vkenv/nvidia570/lib/libnvidia-ml.so.570.172.08}"

export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4
export CUDA_VISIBLE_DEVICES="${ONF_GPU:-0}"
export PYTHONNOUSERSITE=1
export PYTHONPATH="$REPO/stablevla:$REPO/evals/libero_plus/harness:$REPO/src${PYTHONPATH:+:$PYTHONPATH}"
export MAGICK_HOME="$CONDA_PREFIX_DIR"
export LD_LIBRARY_PATH="$CONDA_PREFIX_DIR/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
# torch's caching allocator calls nvmlInit_v2 whenever it grows its pool and hard-asserts on failure;
# this box's userspace NVML (580) does not match its kernel module (570).
[ -f "$NVML_PRELOAD" ] && export LD_PRELOAD="$NVML_PRELOAD${LD_PRELOAD:+:$LD_PRELOAD}"

cd "$REPO"
SERVER_LOG="$REPO/logs/_drivers/policy_server_${SUITE}_${PORT}.log"
SERVER_PID=""

port_open() { (exec 3<>"/dev/tcp/127.0.0.1/$PORT") 2>/dev/null; }

cleanup() {
    [ -n "$SERVER_PID" ] || return 0
    echo "[$(date -u +%FT%TZ)] stopping policy server pid $SERVER_PID"
    kill "$SERVER_PID" 2>/dev/null
    sleep 10
    kill -9 "$SERVER_PID" 2>/dev/null
    pkill -f "stablevla_policy_server.py.*--port $PORT" 2>/dev/null
}
trap cleanup EXIT INT TERM

if port_open; then
    echo "[$(date -u +%FT%TZ)] reusing the policy server already listening on 127.0.0.1:$PORT"
else
    CKPT="$("$PY" -c "import sys; sys.path.insert(0, '$REPO/scripts'); \
import cache_policy_actions as c; print(c.suite_checkpoint('$SUITE'))")" || exit 1
    echo "[$(date -u +%FT%TZ)] starting policy server: $CKPT -> port $PORT (cuda 0), log $SERVER_LOG"
    setsid "$PY" evals/stablevla/stablevla_policy_server.py \
        --ckpt_path "$CKPT" --port "$PORT" --cuda 0 --use_bf16 >>"$SERVER_LOG" 2>&1 </dev/null &
    SERVER_PID=$!
    for _ in $(seq 600); do
        port_open && break
        kill -0 "$SERVER_PID" 2>/dev/null || { echo "!!! server died; see $SERVER_LOG"; exit 1; }
        sleep 2
    done
    port_open || { echo "!!! server not ready after 1200s; see $SERVER_LOG"; exit 1; }
fi

echo "[$(date -u +%FT%TZ)] === caching $SUITE policy actions on port $PORT"
"$PY" scripts/cache_policy_actions.py "$SUITE" --port "$PORT" "$@"
RC=$?
echo "[$(date -u +%FT%TZ)] === exit $RC"
exit $RC