File size: 2,478 Bytes
23e02ba | 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 | #!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/workspace/mp_20_pxrdnet}
ENV_ACTIVATE=${ENV_ACTIVATE:-/workspace/miniforge/bin/activate}
ENV_NAME=${ENV_NAME:-pxrdnet}
MODEL_PATH=${MODEL_PATH:-${ROOT}/hydra/singlerun/2026-07-20/pxrdgen_raw512_run02}
OUTPUT_ROOT=${OUTPUT_ROOT:-${ROOT}/paper_results_pxrdgen_match_only/cpu_multinode/full_test}
TOTAL_MATERIALS=${TOTAL_MATERIALS:-9046}
CHUNK_SIZE=${CHUNK_SIZE:?Set CHUNK_SIZE to the best single-node material count.}
NODE_INDEX=${NODE_INDEX:-${SLURM_ARRAY_TASK_ID:-0}}
FIRST_IDX=$((NODE_INDEX * CHUNK_SIZE))
if (( FIRST_IDX >= TOTAL_MATERIALS )); then
echo "NODE_INDEX=${NODE_INDEX} starts at ${FIRST_IDX}, beyond TOTAL_MATERIALS=${TOTAL_MATERIALS}; nothing to do."
exit 0
fi
REMAINING=$((TOTAL_MATERIALS - FIRST_IDX))
if (( REMAINING < CHUNK_SIZE )); then
NUM_MATERIALS=${REMAINING}
else
NUM_MATERIALS=${CHUNK_SIZE}
fi
OUTPUT_DIR=${OUTPUT_DIR:-${OUTPUT_ROOT}/node_$(printf "%05d" "${NODE_INDEX}")}
CPU_WORKERS=${CPU_WORKERS:-0}
TORCH_THREADS_PER_WORKER=${TORCH_THREADS_PER_WORKER:-0}
NUM_STARTING_POINTS=${NUM_STARTING_POINTS:-1}
NUM_GRADIENT_STEPS=${NUM_GRADIENT_STEPS:-5000}
N_STEP_EACH=${N_STEP_EACH:-100}
PROGRESS_LOG_INTERVAL=${PROGRESS_LOG_INTERVAL:-500}
PROGRESS_MININTERVAL=${PROGRESS_MININTERVAL:-10}
source "${ENV_ACTIVATE}" "${ENV_NAME}"
cd "${ROOT}"
export PROJECT_ROOT="${ROOT}"
export HYDRA_JOBS="${ROOT}/hydra"
export WABDB_DIR="${ROOT}/wabdb"
export PYTHONPATH="${ROOT}:${ROOT}/scripts:${PYTHONPATH:-}"
export WANDB_MODE=disabled
export CUDA_VISIBLE_DEVICES=""
export PIP_CACHE_DIR="${PIP_CACHE_DIR:-/workspace/.cache/pip}"
export TMPDIR="${TMPDIR:-/workspace/tmp}"
mkdir -p "${OUTPUT_DIR}"
echo "Running node chunk:"
echo " NODE_INDEX=${NODE_INDEX}"
echo " FIRST_IDX=${FIRST_IDX}"
echo " NUM_MATERIALS=${NUM_MATERIALS}"
echo " CHUNK_SIZE=${CHUNK_SIZE}"
echo " OUTPUT_DIR=${OUTPUT_DIR}"
python -W ignore "${ROOT}/cpu_generation_benchmark/run_cpu_generation_benchmark.py" \
--model-path "${MODEL_PATH}" \
--output-dir "${OUTPUT_DIR}" \
--first-idx "${FIRST_IDX}" \
--num-materials "${NUM_MATERIALS}" \
--num-starting-points "${NUM_STARTING_POINTS}" \
--workers "${CPU_WORKERS}" \
--torch-threads-per-worker "${TORCH_THREADS_PER_WORKER}" \
--num-gradient-steps "${NUM_GRADIENT_STEPS}" \
--n-step-each "${N_STEP_EACH}" \
--progress-log-interval "${PROGRESS_LOG_INTERVAL}" \
--progress-mininterval "${PROGRESS_MININTERVAL}" \
2>&1 | tee "${OUTPUT_DIR}/benchmark.log"
|