SignerX/SignVerse-2M

1#!/usr/bin/env bash2#SBATCH --job-name=dwpose3#SBATCH --nodes=14#SBATCH --ntasks=15#SBATCH --cpus-per-task=86#SBATCH --gres=gpu:17#SBATCH --mem=32G8#SBATCH --time=24:00:009#SBATCH --output=%x_%A_%a.out10#SBATCH --error=%x_%A_%a.err11 12set -euo pipefail13 14ROOT_DIR="${ROOT_DIR:-/home/sf895/SignVerse-2M}"15RUNTIME_ROOT="${RUNTIME_ROOT:-/home/sf895/SignVerse-2M-runtime}"16STATE_ROOT="${STATE_ROOT:-/home/sf895/SignVerse-2M-runtime}"17CONDA_SH="${CONDA_SH:-/home/sf895/miniconda3/etc/profile.d/conda.sh}"18CONDA_ENV="${CONDA_ENV:-signx2}"19RAW_VIDEO_DIR="${RAW_VIDEO_DIR:-$RUNTIME_ROOT/raw_video}"20SCRATCH_RAW_VIDEO_DIR="${SCRATCH_RAW_VIDEO_DIR:-/scratch/$USER/SignVerse-2M-runtime/raw_video}"21DATASET_DIR="${DATASET_DIR:-$RUNTIME_ROOT/dataset}"22SCRATCH_DATASET_DIR="${SCRATCH_DATASET_DIR:-/scratch/$USER/SignVerse-2M-runtime/dataset}"23STATS_NPZ="${STATS_NPZ:-$RUNTIME_ROOT/stats.npz}"24STATUS_JOURNAL_PATH="${STATUS_JOURNAL_PATH:-$RUNTIME_ROOT/upload_status_journal.jsonl}"25PIPELINE02="${PIPELINE02:-$ROOT_DIR/scripts/pipeline02_extract_dwpose_from_video.py}"26FPS="${FPS:-24}"27TMP_ROOT="${TMP_ROOT:-${SLURM_TMPDIR:-/tmp}}"28FORCE_PROCESS="${FORCE_PROCESS:-0}"29DELETE_SOURCE_ON_SUCCESS="${DELETE_SOURCE_ON_SUCCESS:-0}"30CLAIM_DIR="${CLAIM_DIR:-$STATE_ROOT/slurm/state/claims}"31RETRY_DIR="${RETRY_DIR:-$STATE_ROOT/slurm/state/gpu_init_retries}"32MAX_GPU_INIT_RETRIES="${MAX_GPU_INIT_RETRIES:-3}"33OPTIMIZED_MODE="${OPTIMIZED_MODE:-1}"34OPTIMIZED_PROVIDER="${OPTIMIZED_PROVIDER:-cuda}"35OPTIMIZED_FRAME_BATCH_SIZE="${OPTIMIZED_FRAME_BATCH_SIZE:-8}"36OPTIMIZED_DETECT_RESOLUTION="${OPTIMIZED_DETECT_RESOLUTION:-512}"37OPTIMIZED_FRAME_STRIDE="${OPTIMIZED_FRAME_STRIDE:-1}"38OPTIMIZED_IO_BINDING="${OPTIMIZED_IO_BINDING:-1}"39OPTIMIZED_GPU_DETECTOR_POSTPROCESS="${OPTIMIZED_GPU_DETECTOR_POSTPROCESS:-1}"40OPTIMIZED_GPU_POSE_PREPROCESS="${OPTIMIZED_GPU_POSE_PREPROCESS:-0}"41VIDEOS_PER_JOB="${VIDEOS_PER_JOB:-5}"42 43MANIFEST="${MANIFEST:-${1:-}}"44if [[ -z "$MANIFEST" ]]; then45  echo "MANIFEST is required (env var or first positional arg)." >&246  exit 147fi48if [[ ! -f "$MANIFEST" ]]; then49  echo "Manifest not found: $MANIFEST" >&250  exit 151fi52if [[ -z "${SLURM_ARRAY_TASK_ID:-}" ]]; then53  echo "SLURM_ARRAY_TASK_ID is required." >&254  exit 155fi56if [[ ! -f "$CONDA_SH" ]]; then57  echo "Missing conda init script: $CONDA_SH" >&258  exit 159fi60 61mapfile -t ALL_VIDEO_IDS < <(sed -n "$((SLURM_ARRAY_TASK_ID * VIDEOS_PER_JOB + 1)),$(((SLURM_ARRAY_TASK_ID + 1) * VIDEOS_PER_JOB))p" "$MANIFEST")62if [[ "${#ALL_VIDEO_IDS[@]}" -eq 0 ]]; then63  echo "No video ids found for task index ${SLURM_ARRAY_TASK_ID} in manifest $MANIFEST" >&264  exit 165fi66 67mkdir -p "$CLAIM_DIR" "$RETRY_DIR"68cleanup_claims() {69  local video_id70  for video_id in "${ALL_VIDEO_IDS[@]}"; do71    rm -f "$CLAIM_DIR/${video_id}.claim"72  done73}74trap cleanup_claims EXIT75 76export OMP_NUM_THREADS="${SLURM_CPUS_PER_TASK:-1}"77export MKL_NUM_THREADS="${SLURM_CPUS_PER_TASK:-1}"78 79echo "[$(date '+%F %T')] job=${SLURM_JOB_ID:-na} task=${SLURM_ARRAY_TASK_ID} host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES:-unset} videos=${#ALL_VIDEO_IDS[@]} first_video=${ALL_VIDEO_IDS[0]}"80 81# shellcheck disable=SC109082source "$CONDA_SH"83CONDA_ENV_PREFIX="$(conda env list | awk '$1 == env {print $NF}' env="$CONDA_ENV")"84if [[ -z "$CONDA_ENV_PREFIX" ]]; then85  echo "Unable to resolve conda env prefix for $CONDA_ENV" >&286  exit 187fi88export LD_LIBRARY_PATH="$CONDA_ENV_PREFIX/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"89export PATH="$CONDA_ENV_PREFIX/bin:$PATH"90LIBSTDCXX_PATH="$CONDA_ENV_PREFIX/lib/libstdc++.so.6"91LIBGCC_PATH="$CONDA_ENV_PREFIX/lib/libgcc_s.so.1"92LD_PRELOAD_VALUE=""93if [[ -f "$LIBSTDCXX_PATH" ]]; then94  LD_PRELOAD_VALUE="$LIBSTDCXX_PATH"95fi96if [[ -f "$LIBGCC_PATH" ]]; then97  LD_PRELOAD_VALUE="${LD_PRELOAD_VALUE:+$LD_PRELOAD_VALUE:}$LIBGCC_PATH"98fi99export LD_PRELOAD="${LD_PRELOAD_VALUE}${LD_PRELOAD:+:$LD_PRELOAD}"100echo "Using conda env prefix=$CONDA_ENV_PREFIX"101echo "Using LD_PRELOAD=$LD_PRELOAD"102 103gpu_init_retry_state() {104  local video_id="$1"105  local attempts="$2"106  local host_name="$3"107  local gpu_id="$4"108  local reason="$5"109  local retry_state_path="$RETRY_DIR/${video_id}.state"110  cat > "$retry_state_path" <<STATE111attempts=$attempts112host=$host_name113gpu=$gpu_id114updated_at=$(date '+%F %T')115reason=$reason116STATE117}118 119get_retry_attempts() {120  local video_id="$1"121  local retry_state_path="$RETRY_DIR/${video_id}.state"122  if [[ -f "$retry_state_path" ]]; then123    awk -F'=' '/^attempts=/{print $2}' "$retry_state_path" | tail -n 1124  else125    echo 0126  fi127}128 129mark_retry_exhausted() {130  local video_id="$1"131  local attempts="$2"132  local reason="$3"133  python - <<PY134from pathlib import Path135import sys, time136root = Path(r"$ROOT_DIR")137stats_path = Path(r"$STATS_NPZ")138video_id = r"$video_id"139attempts = int(r"$attempts")140reason = r'''$reason'''141sys.path.insert(0, str(root))142from utils.stats_npz import update_video_stats_best_effort143journal_path = Path(r"$STATUS_JOURNAL_PATH")144update_video_stats_best_effort(145    stats_path,146    journal_path,147    video_id,148    process_status="skipped",149    last_error=f"gpu_init_retry_exhausted after {attempts} attempts: {reason}",150    updated_at=time.strftime("%Y-%m-%d %H:%M:%S"),151)152PY153  rm -f "$RAW_VIDEO_DIR/$video_id.mp4" "$RAW_VIDEO_DIR/$video_id.mkv" "$RAW_VIDEO_DIR/$video_id.webm" "$RAW_VIDEO_DIR/$video_id.mov"154  rm -f "$SCRATCH_RAW_VIDEO_DIR/$video_id.mp4" "$SCRATCH_RAW_VIDEO_DIR/$video_id.mkv" "$SCRATCH_RAW_VIDEO_DIR/$video_id.webm" "$SCRATCH_RAW_VIDEO_DIR/$video_id.mov"155}156 157should_retry_gpu_init_failure() {158  local log_path="$1"159  local npz_dir="$2"160  local npz_count="0"161  if [[ -d "$npz_dir" ]]; then162    npz_count="$(find "$npz_dir" -maxdepth 1 -name '*.npz' | wc -l | tr -d '[:space:]')"163  fi164  if [[ "$npz_count" != "0" ]]; then165    return 1166  fi167  grep -Eiq 'CUDA failure 2: out of memory|Failed to create CUDAExecutionProvider|libcudnn\.so\.8|CUDA is not available|CUDAExecutionProvider|onnxruntime' "$log_path"168}169 170env_cmd=(env171  "LD_LIBRARY_PATH=$LD_LIBRARY_PATH"172  "LD_PRELOAD=$LD_PRELOAD"173  "PATH=$PATH"174  "CONDA_NO_PLUGINS=true"175)176 177cmd_base=(python -u "$PIPELINE02"178  --raw-video-dir "$RAW_VIDEO_DIR"179  --scratch-raw-video-dir "$SCRATCH_RAW_VIDEO_DIR"180  --dataset-dir "$DATASET_DIR"181  --scratch-dataset-dir "$SCRATCH_DATASET_DIR"182  --stats-npz "$STATS_NPZ"183  --status-journal-path "$STATUS_JOURNAL_PATH"184  --fps "$FPS"185  --workers 1186  --tmp-root "$TMP_ROOT"187)188if [[ "$OPTIMIZED_MODE" == "1" ]]; then189  cmd_base+=(--optimized-mode)190  cmd_base+=(--optimized-provider "$OPTIMIZED_PROVIDER")191  cmd_base+=(--optimized-frame-batch-size "$OPTIMIZED_FRAME_BATCH_SIZE")192  cmd_base+=(--optimized-detect-resolution "$OPTIMIZED_DETECT_RESOLUTION")193  cmd_base+=(--optimized-frame-stride "$OPTIMIZED_FRAME_STRIDE")194  if [[ "$OPTIMIZED_IO_BINDING" == "1" ]]; then195    cmd_base+=(--optimized-io-binding)196  fi197  if [[ "$OPTIMIZED_GPU_DETECTOR_POSTPROCESS" == "1" ]]; then198    cmd_base+=(--optimized-gpu-detector-postprocess)199  fi200  if [[ "$OPTIMIZED_GPU_POSE_PREPROCESS" == "1" ]]; then201    cmd_base+=(--optimized-gpu-pose-preprocess)202  fi203else204  cmd_base+=(--legacy-mode)205fi206if [[ "$FORCE_PROCESS" == "1" ]]; then207  cmd_base+=(--force)208fi209if [[ "$DELETE_SOURCE_ON_SUCCESS" == "1" ]]; then210  cmd_base+=(--delete-source-on-success)211fi212 213overall_status=0214for VIDEO_ID in "${ALL_VIDEO_IDS[@]}"; do215  [[ -z "$VIDEO_ID" ]] && continue216  CLAIM_PATH="$CLAIM_DIR/${VIDEO_ID}.claim"217  RETRY_STATE_PATH="$RETRY_DIR/${VIDEO_ID}.state"218  echo "[$(date '+%F %T')] job=${SLURM_JOB_ID:-na} task=${SLURM_ARRAY_TASK_ID} host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES:-unset} video_id=$VIDEO_ID"219 220  TMP_LOG="$(mktemp "${TMP_ROOT%/}/dwpose_${VIDEO_ID}_XXXX.log")"221  cmd=("${cmd_base[@]}" --video-ids="$VIDEO_ID")222  set +e223  "${env_cmd[@]}" conda run -n "$CONDA_ENV" "${cmd[@]}" 2>&1 | tee "$TMP_LOG"224  cmd_status=${PIPESTATUS[0]}225  set -e226 227  if [[ "$cmd_status" -ne 0 ]]; then228    npz_dir="$DATASET_DIR/$VIDEO_ID/npz"229    if [[ ! -d "$npz_dir" && -d "$SCRATCH_DATASET_DIR/$VIDEO_ID/npz" ]]; then230      npz_dir="$SCRATCH_DATASET_DIR/$VIDEO_ID/npz"231    fi232    if should_retry_gpu_init_failure "$TMP_LOG" "$npz_dir"; then233      attempts="$(get_retry_attempts "$VIDEO_ID")"234      attempts="$((attempts + 1))"235      last_reason="$(tail -n 80 "$TMP_LOG" | tr '\n' ' ' | sed 's/[[:space:]]\+/ /g' | cut -c1-1200)"236      gpu_init_retry_state "$VIDEO_ID" "$attempts" "$(hostname)" "${CUDA_VISIBLE_DEVICES:-unset}" "$last_reason"237      if [[ "$attempts" -ge "$MAX_GPU_INIT_RETRIES" ]]; then238        echo "GPU init failed on multiple GPUs; marking $VIDEO_ID as skipped after $attempts attempts." >&2239        mark_retry_exhausted "$VIDEO_ID" "$attempts" "$last_reason"240        rm -f "$CLAIM_PATH"241        continue242      fi243      echo "GPU init failure for $VIDEO_ID on host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES:-unset}; retry attempt $attempts/$MAX_GPU_INIT_RETRIES will be resubmitted later." >&2244      overall_status="$cmd_status"245      break246    fi247    overall_status="$cmd_status"248    break249  fi250 251  rm -f "$RETRY_STATE_PATH" "$CLAIM_PATH"252done253 254exit "$overall_status"255