Add script for extracting EasyCom sessions with three workers per GPU. This script manages session assignments, ensures serialized worker launches, and includes error handling for session readiness and logging. Enhances the EasyCom pipeline for improved session management and feature extraction.
This commit is contained in:
Xin Wang
@@ -0,0 +1,247 @@
|
||||
#!/usr/bin/env bash
|
||||
# Extract all 12 EasyCom sessions with three independent workers per GPU.
|
||||
#
|
||||
# Sessions are assigned round-robin across four GPUs:
|
||||
# GPU 0: Session 1, 5, 9
|
||||
# GPU 1: Session 2, 6, 10
|
||||
# GPU 2: Session 3, 7, 11
|
||||
# GPU 3: Session 4, 8, 12
|
||||
#
|
||||
# Worker launches are deliberately serialized. The launcher waits until one
|
||||
# worker has loaded Qwen and produced its first 0.2-second state before starting
|
||||
# the next worker, preventing all model checkpoints from loading at once.
|
||||
set -uo pipefail
|
||||
|
||||
PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}"
|
||||
MODEL="${MODEL:-/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B}"
|
||||
VIDEO_ROOT="${VIDEO_ROOT:-/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed}"
|
||||
TARGET_ROOT="${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame}"
|
||||
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct}"
|
||||
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct}"
|
||||
|
||||
GPUS="${GPUS:-0 1 2 3}"
|
||||
SESSIONS="${SESSIONS:-1 2 3 4 5 6 7 8 9 10 11 12}"
|
||||
PROCESSES_PER_GPU="${PROCESSES_PER_GPU:-3}"
|
||||
MEMORY_SECONDS="${MEMORY_SECONDS:-20}"
|
||||
VIDEO_FPS="${VIDEO_FPS:-2}"
|
||||
FRAME_CACHE_SIZE="${FRAME_CACHE_SIZE:-64}"
|
||||
DTYPE="${DTYPE:-bfloat16}"
|
||||
OMP_THREADS_PER_PROCESS="${OMP_THREADS_PER_PROCESS:-4}"
|
||||
|
||||
# Maximum time allowed for one worker to index media, load Qwen, and produce
|
||||
# its first state. A small gap after readiness further separates CUDA startup.
|
||||
START_TIMEOUT_SECONDS="${START_TIMEOUT_SECONDS:-1800}"
|
||||
START_GAP_SECONDS="${START_GAP_SECONDS:-10}"
|
||||
POLL_SECONDS="${POLL_SECONDS:-5}"
|
||||
|
||||
read -r -a GPU_ARRAY <<<"${GPUS}"
|
||||
read -r -a SESSION_ARRAY <<<"${SESSIONS}"
|
||||
|
||||
if [[ "${#GPU_ARRAY[@]}" -eq 0 ]]; then
|
||||
echo "No GPUs supplied through GPUS" >&2
|
||||
exit 2
|
||||
fi
|
||||
if [[ "${PROCESSES_PER_GPU}" -le 0 ]]; then
|
||||
echo "PROCESSES_PER_GPU must be positive" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
capacity=$((${#GPU_ARRAY[@]} * PROCESSES_PER_GPU))
|
||||
if [[ "${#SESSION_ARRAY[@]}" -gt "${capacity}" ]]; then
|
||||
echo "${#SESSION_ARRAY[@]} sessions exceed capacity ${capacity}" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"
|
||||
cd "${PROJECT_ROOT}"
|
||||
|
||||
declare -a CHILD_PIDS=()
|
||||
declare -a CHILD_NAMES=()
|
||||
declare -a CHILD_PIDFILES=()
|
||||
declare -a EXTERNAL_PIDS=()
|
||||
declare -a EXTERNAL_NAMES=()
|
||||
declare -a EXTERNAL_OUTPUTS=()
|
||||
declare -a EXTERNAL_PIDFILES=()
|
||||
declare -A GPU_JOB_COUNTS=()
|
||||
|
||||
for gpu in "${GPU_ARRAY[@]}"; do
|
||||
GPU_JOB_COUNTS["${gpu}"]=0
|
||||
done
|
||||
|
||||
wait_until_ready() {
|
||||
local pid="$1"
|
||||
local log="$2"
|
||||
local session="$3"
|
||||
local started_at="${SECONDS}"
|
||||
|
||||
while true; do
|
||||
# tqdm writes progress updates with carriage returns; grep -a handles the
|
||||
# resulting log as text. This marker appears after the first model forward.
|
||||
if grep -aq "media=0.2s" "${log}"; then
|
||||
return 0
|
||||
fi
|
||||
if grep -aq "Traceback (most recent call last)" "${log}"; then
|
||||
echo "${session} failed during startup; inspect ${log}" >&2
|
||||
return 1
|
||||
fi
|
||||
if ! kill -0 "${pid}" 2>/dev/null; then
|
||||
echo "${session} exited before producing its first state; inspect ${log}" >&2
|
||||
return 1
|
||||
fi
|
||||
if ((SECONDS - started_at >= START_TIMEOUT_SECONDS)); then
|
||||
echo "${session} startup timed out after ${START_TIMEOUT_SECONDS}s" >&2
|
||||
return 1
|
||||
fi
|
||||
sleep "${POLL_SECONDS}"
|
||||
done
|
||||
}
|
||||
|
||||
startup_status=0
|
||||
for index in "${!SESSION_ARRAY[@]}"; do
|
||||
sid="${SESSION_ARRAY[${index}]}"
|
||||
gpu="${GPU_ARRAY[$((index % ${#GPU_ARRAY[@]}))]}"
|
||||
session="Session_${sid}"
|
||||
session_dir="${VIDEO_ROOT}/${session}"
|
||||
target="${TARGET_ROOT}/${session}.npy"
|
||||
output="${FEATURE_ROOT}/${session}.pt"
|
||||
log="${LOG_ROOT}/${session}.log"
|
||||
pidfile="${LOG_ROOT}/${session}.pid"
|
||||
|
||||
GPU_JOB_COUNTS["${gpu}"]=$((GPU_JOB_COUNTS["${gpu}"] + 1))
|
||||
if [[ "${GPU_JOB_COUNTS[${gpu}]}" -gt "${PROCESSES_PER_GPU}" ]]; then
|
||||
echo "GPU ${gpu} was assigned too many workers" >&2
|
||||
startup_status=1
|
||||
break
|
||||
fi
|
||||
if [[ ! -d "${session_dir}" ]]; then
|
||||
echo "Missing session directory: ${session_dir}" >&2
|
||||
startup_status=1
|
||||
break
|
||||
fi
|
||||
if [[ ! -s "${target}" ]]; then
|
||||
echo "Missing target: ${target}" >&2
|
||||
startup_status=1
|
||||
break
|
||||
fi
|
||||
if [[ -s "${output}" ]]; then
|
||||
echo "Skipping completed ${session}: ${output}"
|
||||
continue
|
||||
fi
|
||||
|
||||
# Also detect a worker launched manually before this manager. This is useful
|
||||
# when Session_1 is already running on GPU 0. Its round-robin slot still
|
||||
# counts toward that GPU's three-worker limit.
|
||||
running_pid=""
|
||||
if command -v pgrep >/dev/null 2>&1; then
|
||||
running_pid=$(pgrep -f \
|
||||
"streaming_turn_model_directly.py.*${session_dir}" | head -n 1 || true)
|
||||
fi
|
||||
if [[ -n "${running_pid}" ]] && kill -0 "${running_pid}" 2>/dev/null; then
|
||||
echo "Skipping already-running ${session}: PID=${running_pid}"
|
||||
echo "${running_pid}" >"${pidfile}"
|
||||
EXTERNAL_PIDS+=("${running_pid}")
|
||||
EXTERNAL_NAMES+=("${session}")
|
||||
EXTERNAL_OUTPUTS+=("${output}")
|
||||
EXTERNAL_PIDFILES+=("${pidfile}")
|
||||
if [[ -f "${log}" ]]; then
|
||||
echo " Confirming the existing worker reached its first state..."
|
||||
if ! wait_until_ready "${running_pid}" "${log}" "${session}"; then
|
||||
startup_status=1
|
||||
break
|
||||
fi
|
||||
else
|
||||
echo " WARNING: no log is available to confirm startup; waiting " \
|
||||
"${START_GAP_SECONDS}s before the next load" >&2
|
||||
fi
|
||||
sleep "${START_GAP_SECONDS}"
|
||||
continue
|
||||
fi
|
||||
if [[ -s "${pidfile}" ]]; then
|
||||
existing_pid=$(<"${pidfile}")
|
||||
if kill -0 "${existing_pid}" 2>/dev/null; then
|
||||
echo "Skipping already-running ${session}: PID=${existing_pid}"
|
||||
continue
|
||||
fi
|
||||
fi
|
||||
|
||||
: >"${log}"
|
||||
echo "Starting ${session} on physical GPU ${gpu} " \
|
||||
"(${GPU_JOB_COUNTS[${gpu}]}/${PROCESSES_PER_GPU} workers)"
|
||||
|
||||
env \
|
||||
CUDA_VISIBLE_DEVICES="${gpu}" \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
|
||||
OMP_NUM_THREADS="${OMP_THREADS_PER_PROCESS}" \
|
||||
python streaming_turn_model_directly.py \
|
||||
--model-id "${MODEL}" \
|
||||
--easycom-session-dir "${session_dir}" \
|
||||
--target "${target}" \
|
||||
--use-audio-in-video \
|
||||
--mode open-stream-replay \
|
||||
--memory-seconds "${MEMORY_SECONDS}" \
|
||||
--video-fps "${VIDEO_FPS}" \
|
||||
--frame-cache-size "${FRAME_CACHE_SIZE}" \
|
||||
--dtype "${DTYPE}" \
|
||||
--flash-attention-2 \
|
||||
--device cuda:0 \
|
||||
--output "${output}" \
|
||||
>"${log}" 2>&1 &
|
||||
|
||||
pid=$!
|
||||
echo "${pid}" >"${pidfile}"
|
||||
CHILD_PIDS+=("${pid}")
|
||||
CHILD_NAMES+=("${session}")
|
||||
CHILD_PIDFILES+=("${pidfile}")
|
||||
echo " PID=${pid}; waiting for first state before next model load..."
|
||||
|
||||
if ! wait_until_ready "${pid}" "${log}" "${session}"; then
|
||||
startup_status=1
|
||||
tail -n 40 "${log}" >&2
|
||||
break
|
||||
fi
|
||||
echo " ${session} is running; next worker starts in ${START_GAP_SECONDS}s"
|
||||
sleep "${START_GAP_SECONDS}"
|
||||
done
|
||||
|
||||
echo "Launcher finished starting workers; waiting for child processes."
|
||||
final_status="${startup_status}"
|
||||
for index in "${!CHILD_PIDS[@]}"; do
|
||||
pid="${CHILD_PIDS[${index}]}"
|
||||
session="${CHILD_NAMES[${index}]}"
|
||||
pidfile="${CHILD_PIDFILES[${index}]}"
|
||||
if wait "${pid}"; then
|
||||
echo "Completed ${session}"
|
||||
else
|
||||
echo "Failed ${session}; inspect ${LOG_ROOT}/${session}.log" >&2
|
||||
final_status=1
|
||||
fi
|
||||
rm -f "${pidfile}"
|
||||
done
|
||||
|
||||
# Non-child workers cannot be handled by bash's wait builtin. Poll any workers
|
||||
# that were already running when this launcher started, then validate that they
|
||||
# produced their expected output before reporting global success.
|
||||
for index in "${!EXTERNAL_PIDS[@]}"; do
|
||||
pid="${EXTERNAL_PIDS[${index}]}"
|
||||
session="${EXTERNAL_NAMES[${index}]}"
|
||||
output="${EXTERNAL_OUTPUTS[${index}]}"
|
||||
pidfile="${EXTERNAL_PIDFILES[${index}]}"
|
||||
while kill -0 "${pid}" 2>/dev/null; do
|
||||
sleep 30
|
||||
done
|
||||
if [[ -s "${output}" ]]; then
|
||||
echo "Completed pre-existing ${session}"
|
||||
else
|
||||
echo "Pre-existing ${session} exited without ${output}" >&2
|
||||
final_status=1
|
||||
fi
|
||||
rm -f "${pidfile}"
|
||||
done
|
||||
|
||||
if [[ "${final_status}" -ne 0 ]]; then
|
||||
echo "One or more workers failed, or worker startup was interrupted." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "All selected EasyCom sessions completed: ${FEATURE_ROOT}"
|
||||
Reference in New Issue
Block a user