mp_20_pxrdnet / cpu_multinode_generation /slurm_array_example.sh
2090741942justin's picture
Add CPU multi-node generation orchestration
23e02ba
Raw
History Blame Contribute Delete
1.03 kB
#!/usr/bin/env bash
#SBATCH --job-name=pxrdnet-cpu-gen
#SBATCH --output=/workspace/mp_20_pxrdnet/paper_results_pxrdgen_match_only/cpu_multinode/full_test/slurm_%A_%a.out
#SBATCH --error=/workspace/mp_20_pxrdnet/paper_results_pxrdgen_match_only/cpu_multinode/full_test/slurm_%A_%a.err
#SBATCH --mem=0
#SBATCH --time=48:00:00
#
# Submit after choosing CHUNK_SIZE:
# CHUNK_SIZE=<BEST_NUM_MATERIALS> sbatch --array=0-<NODE_COUNT_MINUS_1> cpu_multinode_generation/slurm_array_example.sh
set -euo pipefail
ROOT=${ROOT:-/workspace/mp_20_pxrdnet}
CHUNK_SIZE=${CHUNK_SIZE:?Set CHUNK_SIZE before sbatch, e.g. CHUNK_SIZE=48.}
NODE_INDEX=${SLURM_ARRAY_TASK_ID}
CPU_WORKERS=${CPU_WORKERS:-0}
TORCH_THREADS_PER_WORKER=${TORCH_THREADS_PER_WORKER:-0}
TOTAL_MATERIALS=${TOTAL_MATERIALS:-9046}
OUTPUT_ROOT=${OUTPUT_ROOT:-${ROOT}/paper_results_pxrdgen_match_only/cpu_multinode/full_test}
export ROOT CHUNK_SIZE NODE_INDEX CPU_WORKERS TORCH_THREADS_PER_WORKER TOTAL_MATERIALS OUTPUT_ROOT
bash "${ROOT}/cpu_multinode_generation/run_node_chunk.sh"