#!/usr/bin/env bash
#SBATCH --job-name=pxrdnet-cpu-wave
#SBATCH --output=/workspace/mp_20_pxrdnet/paper_results_pxrdgen_match_only/cpu_multinode/full_test/slurm_wave_%A_%a.out
#SBATCH --error=/workspace/mp_20_pxrdnet/paper_results_pxrdgen_match_only/cpu_multinode/full_test/slurm_wave_%A_%a.err
#SBATCH --mem=0
#SBATCH --time=48:00:00
#
# Submit one wave at a time:
#   CHUNK_SIZE=<BEST> ACTUAL_NODES=<REAL_NODE_COUNT> WAVE_INDEX=0 \
#     sbatch --array=0-<REAL_NODE_COUNT_MINUS_1> cpu_multinode_generation/slurm_wave_example.sh
#
# Then submit WAVE_INDEX=1, WAVE_INDEX=2, ... until all chunks are complete.

set -euo pipefail

ROOT=${ROOT:-/workspace/mp_20_pxrdnet}
CHUNK_SIZE=${CHUNK_SIZE:?Set CHUNK_SIZE before sbatch.}
ACTUAL_NODES=${ACTUAL_NODES:?Set ACTUAL_NODES before sbatch.}
WAVE_INDEX=${WAVE_INDEX:?Set WAVE_INDEX before sbatch.}
NODE_RANK=${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 ACTUAL_NODES WAVE_INDEX NODE_RANK CPU_WORKERS TORCH_THREADS_PER_WORKER TOTAL_MATERIALS OUTPUT_ROOT

bash "${ROOT}/cpu_multinode_generation/run_wave_node.sh"