Skip to content

Commit f09105a

Browse files
xuelunshenxuelunshen
committed
Update process_videos.sh
Fix the BUG that causes OOM when multiple tasks run on the same GPU.
1 parent 7b04e5e commit f09105a

1 file changed

Lines changed: 128 additions & 50 deletions

File tree

‎process_videos.sh‎

Lines changed: 128 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -13,66 +13,144 @@ if [[ ! -f "$1" || "${1##*.}" != "txt" ]]; then
1313
exit 1
1414
fi
1515

16-
mkdir logs
16+
mkdir -p logs
1717
mkdir -p data/ZeroMatch/pseudo
1818
mkdir -p data/ZeroMatch/video_1080p
1919

20-
function select_gpu() {
21-
while true; do
22-
readarray -t total_memory < <(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits)
23-
readarray -t memory_free < <(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits)
20+
# Create GPU lock file directory
21+
GPU_LOCK_DIR="/tmp/gpu_locks"
22+
mkdir -p "$GPU_LOCK_DIR"
2423

25-
for i in "${!memory_free[@]}"; do
26-
local free_percent=$(awk -v free="${memory_free[$i]}" -v total="${total_memory[$i]}" 'BEGIN{print (free/total)*100}')
27-
if (( $(awk -v fp="$free_percent" -v tp=95 'BEGIN{print (fp >= tp)}') )); then
28-
echo $i
29-
return
30-
fi
31-
done
24+
# Get number of available GPUs
25+
NUM_GPUS=$(nvidia-smi --list-gpus | wc -l)
3226

33-
sleep 30
34-
done
27+
function cleanup_locks() {
28+
rm -f "$GPU_LOCK_DIR"/*.lock
29+
exit
3530
}
3631

37-
while IFS= read -r VIDEO_ID
38-
do
39-
VIDEO_ID=$(echo "$VIDEO_ID" | tr -d '[:space:]')
32+
# Set signal handler to cleanup lock files when script exits
33+
trap cleanup_locks EXIT INT TERM
4034

41-
output_file="./data/ZeroMatch/video_1080p/${VIDEO_ID}.mp4"
42-
if [ ! -f "$output_file" ]; then
43-
yt-dlp -f 'bv*[ext=mp4][height=1080]+ba[ext=m4a]/b[ext=mp4][height=1080]' -S "height, fps" "https://www.youtube.com/watch?v=$VIDEO_ID" -o "$output_file"
44-
else
45-
echo "Video file $output_file already exists. Skipping download."
46-
fi
35+
function select_gpu_safe() {
36+
local max_attempts=60 # Maximum 30 minutes of attempts
37+
local attempt=0
38+
39+
while [ $attempt -lt $max_attempts ]; do
40+
# Check memory status of all GPUs
41+
readarray -t total_memory < <(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits)
42+
readarray -t memory_free < <(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits)
43+
44+
for i in "${!memory_free[@]}"; do
45+
local lock_file="$GPU_LOCK_DIR/gpu_${i}.lock"
46+
47+
# Use file lock to avoid race conditions
48+
if (exec 200>"$lock_file"; flock -n 200); then
49+
local free_percent=$(awk -v free="${memory_free[$i]}" -v total="${total_memory[$i]}" 'BEGIN{print (free/total)*100}')
50+
51+
# Lower threshold to 80%, more practical
52+
if (( $(awk -v fp="$free_percent" -v tp=80 'BEGIN{print (fp >= tp)}') )); then
53+
# Create lock file and write process info
54+
echo "$$:$(date)" > "$lock_file"
55+
echo $i
56+
return 0
57+
fi
58+
fi
59+
done
60+
61+
attempt=$((attempt + 1))
62+
echo "All GPUs are busy, waiting 30 seconds... (attempt $attempt/$max_attempts)" >&2
63+
sleep 30
64+
done
65+
66+
echo "Error: GPU wait timeout" >&2
67+
return 1
68+
}
4769

48-
printf "|%s|%s|%s|%s|%s|%s|\n" "======================" "==============" "============" "======" "========" "====="
49-
printf "| %-20s | %-12s | %-10s | %-4s | %-6s | %-3s |\n" "Timestamp" "Video ID" "Method" "Skip" "Resize" "GPU"
50-
printf "|%s|%s|%s|%s|%s|%s|\n" "----------------------" "--------------" "------------" "------" "--------" "-----"
70+
function release_gpu() {
71+
local gpu_id=$1
72+
local lock_file="$GPU_LOCK_DIR/gpu_${gpu_id}.lock"
73+
rm -f "$lock_file"
74+
}
5175

52-
for skip in 0 1 2
53-
do
54-
for method in GIM_DKM GIM_LOFTR GIM_GLUE SIFT
55-
do
56-
gpu=$(select_gpu)
57-
logstamp=$(date +'%Y%m%d_%H%M%S')
58-
timestamp=$(date +"%Y-%m-%d %H:%M:%S")
59-
printf "| %-20s | %-12s | %-10s | %-4s | %-6s | %-3s |\n" "$timestamp" "$VIDEO_ID" "$method" "$skip" "No" "$gpu"
60-
python video_preprocessor.py --gpu=$gpu --scene_name="$VIDEO_ID" --method=$method --skip=$skip > "logs/${VIDEO_ID}_${method}_skip${skip}_${logstamp}.log" 2>&1 &
61-
sleep 30
62-
done
63-
done
76+
# Cleanup function after task completion
77+
function run_task_with_cleanup() {
78+
local gpu=$1
79+
local cmd="$2"
80+
81+
# Run task
82+
eval "$cmd"
83+
local exit_code=$?
84+
85+
# Release GPU lock
86+
release_gpu "$gpu"
87+
88+
return $exit_code
89+
}
6490

65-
for skip in 0 1 2
66-
do
67-
for method in GIM_DKM GIM_LOFTR GIM_GLUE SIFT
68-
do
69-
gpu=$(select_gpu)
70-
logstamp=$(date +'%Y%m%d_%H%M%S')
71-
timestamp=$(date +"%Y-%m-%d %H:%M:%S")
72-
printf "| %-20s | %-12s | %-10s | %-4s | %-6s | %-3s |\n" "$timestamp" "$VIDEO_ID" "$method" "$skip" "Yes" "$gpu"
73-
python video_preprocessor.py --gpu=$gpu --scene_name="$VIDEO_ID" --method=$method --skip=$skip --resize > "logs/${VIDEO_ID}_${method}_skip${skip}_resize_${logstamp}.log" 2>&1 &
74-
sleep 30
91+
while IFS= read -r VIDEO_ID
92+
do
93+
VIDEO_ID=$(echo "$VIDEO_ID" | tr -d '[:space:]')
94+
95+
output_file="./data/ZeroMatch/video_1080p/${VIDEO_ID}.mp4"
96+
if [ ! -f "$output_file" ]; then
97+
yt-dlp -f 'bv*[ext=mp4][height=1080]+ba[ext=m4a]/b[ext=mp4][height=1080]' -S "height, fps" "https://www.youtube.com/watch?v=$VIDEO_ID" -o "$output_file"
98+
else
99+
echo "Video file $output_file already exists. Skipping download."
100+
fi
101+
102+
printf "|%s|%s|%s|%s|%s|%s|\n" "======================" "==============" "============" "======" "========" "====="
103+
printf "| %-20s | %-12s | %-10s | %-4s | %-6s | %-3s |\n" "Timestamp" "Video ID" "Method" "Skip" "Resize" "GPU"
104+
printf "|%s|%s|%s|%s|%s|%s|\n" "----------------------" "--------------" "------------" "------" "--------" "-----"
105+
106+
# Collect all tasks into array
107+
declare -a tasks=()
108+
109+
# First round: no resize
110+
for skip in 0 1 2; do
111+
for method in GIM_DKM GIM_LOFTR GIM_GLUE SIFT; do
112+
tasks+=("$VIDEO_ID|$method|$skip|No")
113+
done
75114
done
76-
done
77-
115+
116+
# Second round: with resize
117+
for skip in 0 1 2; do
118+
for method in GIM_DKM GIM_LOFTR GIM_GLUE SIFT; do
119+
tasks+=("$VIDEO_ID|$method|$skip|Yes")
120+
done
121+
done
122+
123+
# Execute tasks in order
124+
for task in "${tasks[@]}"; do
125+
IFS='|' read -r vid method skip resize <<< "$task"
126+
127+
# Safely acquire GPU
128+
gpu=$(select_gpu_safe)
129+
if [ $? -ne 0 ]; then
130+
echo "Error: Unable to acquire available GPU, skipping task $task" >&2
131+
continue
132+
fi
133+
134+
logstamp=$(date +'%Y%m%d_%H%M%S')
135+
timestamp=$(date +"%Y-%m-%d %H:%M:%S")
136+
printf "| %-20s | %-12s | %-10s | %-4s | %-6s | %-3s |\n" "$timestamp" "$vid" "$method" "$skip" "$resize" "$gpu"
137+
138+
if [ "$resize" = "Yes" ]; then
139+
cmd="python video_preprocessor.py --gpu=$gpu --scene_name=\"$vid\" --method=$method --skip=$skip --resize > \"logs/${vid}_${method}_skip${skip}_resize_${logstamp}.log\" 2>&1"
140+
else
141+
cmd="python video_preprocessor.py --gpu=$gpu --scene_name=\"$vid\" --method=$method --skip=$skip > \"logs/${vid}_${method}_skip${skip}_${logstamp}.log\" 2>&1"
142+
fi
143+
144+
# Run task in background, automatically release GPU lock after completion
145+
(run_task_with_cleanup "$gpu" "$cmd") &
146+
147+
# Increase wait time to ensure task fully starts and occupies GPU memory
148+
sleep 60
149+
done
150+
78151
done < "$1"
152+
153+
# Wait for all background tasks to complete
154+
wait
155+
156+
echo "All tasks completed!"

0 commit comments

Comments
 (0)