File size: 2,760 Bytes
cbe0918 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | #!/bin/bash
# Auto-restart wrapper for dataset generation
# Monitors and restarts until target is reached
OUTPUT="${1:-./data/dataset_60k.pt}"
TARGET="${2:-60000}"
LOG_DIR="$(dirname $OUTPUT)"
echo "=============================================="
echo "Auto-restart Dataset Generator"
echo "Output: $OUTPUT"
echo "Target: $TARGET items"
echo "=============================================="
# Load environment
if [ -f .env ]; then
export $(grep -v '^#' .env | xargs)
fi
get_count() {
# Count = base (output file) + new (checkpoint.new)
local new_checkpoint="${OUTPUT}.checkpoint.new"
local base_count=0
local new_count=0
if [ -f "$OUTPUT" ]; then
base_count=$(python3 -c "import torch; d=torch.load('$OUTPUT', map_location='cpu', weights_only=False, mmap=True); print(len(d))" 2>/dev/null || echo "0")
fi
if [ -f "$new_checkpoint" ]; then
new_count=$(python3 -c "import torch; d=torch.load('$new_checkpoint', map_location='cpu', weights_only=False, mmap=True); print(len(d))" 2>/dev/null || echo "0")
fi
echo $((base_count + new_count))
}
attempt=1
while true; do
current=$(get_count)
echo ""
echo "[Attempt $attempt] Current: $current / $TARGET items"
if [ "$current" -ge "$TARGET" ]; then
echo "=============================================="
echo "COMPLETE! $current items generated"
echo "=============================================="
exit 0
fi
remaining=$((TARGET - current))
echo "[Attempt $attempt] Starting generation (need $remaining more)..."
echo ""
# Run the generator with resume flag
python3 datasets/create_dataset.py \
--output "$OUTPUT" \
--count "$TARGET" \
--resume \
2>&1 | tee -a "${LOG_DIR}/auto_restart.log"
exit_code=$?
# Check result
new_count=$(get_count)
echo ""
echo "[Attempt $attempt] Finished with exit code $exit_code"
echo "[Attempt $attempt] Items after run: $new_count"
if [ "$new_count" -ge "$TARGET" ]; then
echo "=============================================="
echo "COMPLETE! $new_count items generated"
echo "=============================================="
exit 0
fi
# If no progress was made, wait longer before retry
if [ "$new_count" -eq "$current" ]; then
echo "[WARN] No progress made, waiting 30s before retry..."
sleep 30
else
echo "[INFO] Progress: $current -> $new_count (+$((new_count - current)))"
echo "[INFO] Restarting in 5s..."
sleep 5
fi
attempt=$((attempt + 1))
# Safety limit
if [ "$attempt" -gt 50 ]; then
echo "[ERROR] Too many attempts (50), giving up"
exit 1
fi
done
|