add retry loop to script

Author: yf225

benchmarks/run_input_shard.sh

@@ -1,6 +1,27 @@
 [[ -z "$RANK_OFFSET" ]] && { echo "Error: RANK_OFFSET is not set"; exit 1; }
 [[ -z "$SHARD" ]] && { echo "Error: SHARD is not set"; exit 1; }
 [[ -z "$WORLD_SIZE" ]] && { echo "Error: WORLD_SIZE is not set"; exit 1; }
-CUDA_VISIBLE_DEVICES=$((RANK_OFFSET+SHARD)) python benchmarks/run.py --input-shard $((SHARD+1))/${WORLD_SIZE} --metrics tflops,gbps,speedup >benchmarks_autotune_$(date +%s)_input_shard_$((SHARD+1))_of_${WORLD_SIZE}.txt 2>&1
+
+# Capture timestamp once for consistent filename
+TIMESTAMP=$(date +%s)
+OUTPUT_FILE="benchmarks_autotune_${TIMESTAMP}_input_shard_$((SHARD+1))_of_${WORLD_SIZE}.txt"
+
+# Retry until success
+attempt=0
+while true; do
+    attempt=$((attempt + 1))
+    echo "Attempt $attempt: Running benchmark for shard $((SHARD+1))/${WORLD_SIZE}..."
+    
+    HELION_FORCE_DISK_CACHE=1 CUDA_VISIBLE_DEVICES=$((RANK_OFFSET+SHARD)) python benchmarks/run.py --input-shard $((SHARD+1))/${WORLD_SIZE} --metrics tflops,gbps,speedup >"$OUTPUT_FILE" 2>&1
+    
+    exit_code=$?
+    if [ $exit_code -eq 0 ]; then
+        echo "Success! Benchmark completed for shard $((SHARD+1))/${WORLD_SIZE}"
+        break
+    else
+        echo "Failed with exit code $exit_code. Retrying..."
+        sleep 5  # Optional: wait 5 seconds before retrying
+    fi
+done
 
 # SHARD=0 RANK_OFFSET=4 WORLD_SIZE=4 bash benchmarks/run_input_shard.sh