PWM / PWM_code /policy_wm_eval /run_gpu_worker.sh

Add files using upload-large-folder tool

dd253d4 verified 3 months ago

2.85 kB

#!/usr/bin/env bash
# One GPU's share of the sweep. Co-hosts the OSCAR server + IsaacSim eval on the
# same GPU. Runs a list of "wm:task" jobs; (re)starts the server only when the WM
# changes (so order jobs by WM to get a single reload).
#
#   run_gpu_worker.sh <GPU> <OSCAR_PORT> <MASTER_PORT> <wm:task> [<wm:task> ...]
#     wm  = lora25 | base       task = open_drawer|open_laptop|pour_balls|push_box|stack_can|unload_cans
set -u
GPU=$1; OSCAR_PORT=$2; MASTER_PORT=$3; shift 3
JOBS=("$@")

HERE=/root/pwm_local/PWM_code/policy_wm_eval
OSCAR=/root/pwm_local/oscar-public
SERVER_PY=/root/pwm_local/PWM_code/policy_rl/rl_sim/act_wm_residual_oscar/workers/wm/oscar_wm_server.py
NUM_EPISODES=${NUM_EPISODES:-30}
OUT_DIR=${OUT_DIR:-$HERE/eval_out}
LOGDIR=${LOGDIR:-$OUT_DIR/logs}
mkdir -p "$LOGDIR"

lora_path() { case "$1" in lora25) echo "finetune/ckpt_6task_lora25/lora_step16000.pt";; base) echo "";; *) echo "??";; esac; }

SERVER_PGID=""
start_server() {   # $1 = wm
  local wm=$1 lora; lora=$(lora_path "$wm")
  echo "[gpu$GPU] start $wm server :$OSCAR_PORT (lora='${lora:-BASE}')" >&2
  setsid bash -c "cd '$OSCAR' && CUDA_VISIBLE_DEVICES=$GPU PYTHONPATH=. HF_HUB_OFFLINE=1 \
    .venv/bin/torchrun --nproc_per_node=1 --master_port=$MASTER_PORT \
    '$SERVER_PY' --lora '$lora' --lora_r 16 --lora_alpha 32 --port $OSCAR_PORT" \
    > "$LOGDIR/server_gpu${GPU}_${wm}.log" 2>&1 &
  SERVER_PGID=$!
  local i
  for i in $(seq 1 80); do
    if curl -s -m 3 "http://localhost:$OSCAR_PORT/health" 2>/dev/null | grep -q model_loaded; then
      echo "[gpu$GPU] $wm ready" >&2; return 0; fi
    kill -0 "$SERVER_PGID" 2>/dev/null || { echo "[gpu$GPU] server proc died" >&2; return 1; }
    sleep 15
  done
  echo "[gpu$GPU] server ready-timeout" >&2; return 1
}
stop_server() {
  [ -n "$SERVER_PGID" ] || return 0
  kill -TERM -"$SERVER_PGID" 2>/dev/null   # negative = whole process group (setsid leader)
  local i
  for i in $(seq 1 15); do kill -0 -"$SERVER_PGID" 2>/dev/null || break; sleep 2; done
  kill -KILL -"$SERVER_PGID" 2>/dev/null   # force any stragglers
  sleep 4                                   # let CUDA memory actually free before reloading
  SERVER_PGID=""
}
trap 'stop_server' EXIT

cur_wm=""
for job in "${JOBS[@]}"; do
  wm=${job%%:*}; task=${job##*:}
  wmtag=$([ "$wm" = lora25 ] && echo lora25_step16000 || echo base)
  if [ "$wm" != "$cur_wm" ]; then
    stop_server
    start_server "$wm" || { echo "[gpu$GPU] ABORT (server $wm failed)" >&2; exit 1; }
    cur_wm=$wm
  fi
  echo "[gpu$GPU] === eval $task in $wmtag ($NUM_EPISODES ep) ===" >&2
  TASK=$task OSCAR_PORT=$OSCAR_PORT WM_TAG=$wmtag GPU=$GPU \
    NUM_EPISODES=$NUM_EPISODES OUT_DIR=$OUT_DIR \
    bash "$HERE/run_eval.sh" > "$LOGDIR/eval_gpu${GPU}_${wmtag}_${task}.log" 2>&1
  echo "[gpu$GPU] finished $task/$wmtag rc=$?" >&2
done
echo "[gpu$GPU] ALL DONE" >&2

	#!/usr/bin/env bash
	# One GPU's share of the sweep. Co-hosts the OSCAR server + IsaacSim eval on the
	# same GPU. Runs a list of "wm:task" jobs; (re)starts the server only when the WM
	# changes (so order jobs by WM to get a single reload).
	#
	# run_gpu_worker.sh <GPU> <OSCAR_PORT> <MASTER_PORT> <wm:task> [<wm:task> ...]
	# wm = lora25 \| base task = open_drawer\|open_laptop\|pour_balls\|push_box\|stack_can\|unload_cans
	set -u
	GPU=$1; OSCAR_PORT=$2; MASTER_PORT=$3; shift 3
	JOBS=("$@")

	HERE=/root/pwm_local/PWM_code/policy_wm_eval
	OSCAR=/root/pwm_local/oscar-public
	SERVER_PY=/root/pwm_local/PWM_code/policy_rl/rl_sim/act_wm_residual_oscar/workers/wm/oscar_wm_server.py
	NUM_EPISODES=${NUM_EPISODES:-30}
	OUT_DIR=${OUT_DIR:-$HERE/eval_out}
	LOGDIR=${LOGDIR:-$OUT_DIR/logs}
	mkdir -p "$LOGDIR"

	lora_path() { case "$1" in lora25) echo "finetune/ckpt_6task_lora25/lora_step16000.pt";; base) echo "";; *) echo "??";; esac; }

	SERVER_PGID=""
	start_server() { # $1 = wm
	local wm=$1 lora; lora=$(lora_path "$wm")
	echo "[gpu$GPU] start $wm server :$OSCAR_PORT (lora='${lora:-BASE}')" >&2
	setsid bash -c "cd '$OSCAR' && CUDA_VISIBLE_DEVICES=$GPU PYTHONPATH=. HF_HUB_OFFLINE=1 \
	.venv/bin/torchrun --nproc_per_node=1 --master_port=$MASTER_PORT \
	'$SERVER_PY' --lora '$lora' --lora_r 16 --lora_alpha 32 --port $OSCAR_PORT" \
	> "$LOGDIR/server_gpu${GPU}_${wm}.log" 2>&1 &
	SERVER_PGID=$!
	local i
	for i in $(seq 1 80); do
	if curl -s -m 3 "http://localhost:$OSCAR_PORT/health" 2>/dev/null \| grep -q model_loaded; then
	echo "[gpu$GPU] $wm ready" >&2; return 0; fi
	kill -0 "$SERVER_PGID" 2>/dev/null \|\| { echo "[gpu$GPU] server proc died" >&2; return 1; }
	sleep 15
	done
	echo "[gpu$GPU] server ready-timeout" >&2; return 1
	}
	stop_server() {
	[ -n "$SERVER_PGID" ] \|\| return 0
	kill -TERM -"$SERVER_PGID" 2>/dev/null # negative = whole process group (setsid leader)
	local i
	for i in $(seq 1 15); do kill -0 -"$SERVER_PGID" 2>/dev/null \|\| break; sleep 2; done
	kill -KILL -"$SERVER_PGID" 2>/dev/null # force any stragglers
	sleep 4 # let CUDA memory actually free before reloading
	SERVER_PGID=""
	}
	trap 'stop_server' EXIT

	cur_wm=""
	for job in "${JOBS[@]}"; do
	wm=${job%%:}; task=${job##:}
	wmtag=$([ "$wm" = lora25 ] && echo lora25_step16000 \|\| echo base)
	if [ "$wm" != "$cur_wm" ]; then
	stop_server
	start_server "$wm" \|\| { echo "[gpu$GPU] ABORT (server $wm failed)" >&2; exit 1; }
	cur_wm=$wm
	fi
	echo "[gpu$GPU] === eval $task in $wmtag ($NUM_EPISODES ep) ===" >&2
	TASK=$task OSCAR_PORT=$OSCAR_PORT WM_TAG=$wmtag GPU=$GPU \
	NUM_EPISODES=$NUM_EPISODES OUT_DIR=$OUT_DIR \
	bash "$HERE/run_eval.sh" > "$LOGDIR/eval_gpu${GPU}_${wmtag}_${task}.log" 2>&1
	echo "[gpu$GPU] finished $task/$wmtag rc=$?" >&2
	done
	echo "[gpu$GPU] ALL DONE" >&2