From e4fb62ba601d048eae1bec373a40e03fa01976dc Mon Sep 17 00:00:00 2001
From: Xin Wang <your-email@example.com>
Date: Mon, 17 Aug 2026 23:43:34 +0800
Subject: [PATCH] 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.

---
 run_extract_easycom_direct_3pergpu.sh | 247 ++++++++++++++++++++++++++
 1 file changed, 247 insertions(+)
 create mode 100644 run_extract_easycom_direct_3pergpu.sh

diff --git a/run_extract_easycom_direct_3pergpu.sh b/run_extract_easycom_direct_3pergpu.sh
new file mode 100644
index 0000000..3358d06
--- /dev/null
+++ b/run_extract_easycom_direct_3pergpu.sh
@@ -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}"