File size: 1,399 Bytes
58258b8 | 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 | #!/bin/bash
# Launch 8 shards on this worker, one per GPU.
# Worker index is derived from $MY_POD_NAME (Arnold pod name ends with worker-N).
#
# Env vars:
# CONFIG — path to the SDG YAML config (default: math53K config)
# LIMIT — optional --limit value for smoke tests (e.g. LIMIT=128)
# NUM_SHARDS — total shards across all workers (default: 32 = 4 nodes × 8 GPU)
#
# Outputs:
# logs/worker_${WORKER_IDX}/shard_NNN.log
set -e
WORKER_IDX=$(echo "$MY_POD_NAME" | grep -oE 'worker-[0-9]+$' | grep -oE '[0-9]+')
if [ -z "$WORKER_IDX" ]; then
echo "ERROR: could not infer worker index from MY_POD_NAME=$MY_POD_NAME" >&2
exit 1
fi
CONFIG="${CONFIG:-sdg/configs/qwen3_4b_openthoughts3_math53K_instill_n8_valredundancy5_round1.yaml}"
NUM_SHARDS="${NUM_SHARDS:-32}"
GPUS_PER_WORKER=8
START=$((WORKER_IDX * GPUS_PER_WORKER))
mkdir -p "logs/worker_${WORKER_IDX}"
LIMIT_FLAG=""
if [ -n "$LIMIT" ]; then LIMIT_FLAG="--limit $LIMIT"; fi
echo "worker $WORKER_IDX launching shards $START..$((START+GPUS_PER_WORKER-1)) (config=$CONFIG, limit=${LIMIT:-none})"
for i in $(seq 0 $((GPUS_PER_WORKER - 1))); do
SHARD_ID=$((START + i))
CUDA_VISIBLE_DEVICES=$i nohup python -m sdg.generate \
--config "$CONFIG" --shard-id "$SHARD_ID" --num-shards "$NUM_SHARDS" $LIMIT_FLAG \
> "logs/worker_${WORKER_IDX}/shard_${SHARD_ID}.log" 2>&1 &
done
echo "launched 8 shards; PIDs: $(jobs -p)"
|