Better benchmark slurm - to fixup

DifferentiableUniverseInitiative · Mar 4, 2024 · 2333fd7 · 2333fd7
1 parent 3ea2eb1
commit 2333fd7
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 302 deletions.
diff --git a/benchmarks/scripts/README.md b/benchmarks/scripts/README.md
@@ -3,17 +3,19 @@
 
 ```bash
 cd $ALL_CCFRWORK/jaxDecomp
-# any python works
-python benchmarks/scripts/the_scripter.py
 ```
 
 this will run everything
 
-# to run specific benchmarks :
+# Run bencmarks on V100-32G:
 
-check `submit_job_from_another_script` in `benchmarks/scripts/the_scripter.py`
+From the root of the jaxDecomp workspace
 
-# to run only Single jobs
+```bash
+sbatch benchmarks/scripts/v100-32G.slurm
+```
+
+# to run only Single jobs interactively
 
 ```bash
 salloc --account=tkc@v100 --nodes=1  --ntasks-per-node=4 --gres=gpu:4 -C a100 --hint=nomultithread --qos=qos_gpu-dev
@@ -22,7 +24,7 @@ salloc --account=tkc@v100 --nodes=1  --ntasks-per-node=4 --gres=gpu:4 -C a100 --
 then
 
 ```bash
-srun python benchmarks/pfft3d.py 1024 2_2 NCCL path_for_csv
-srun python benchmarks/jaxfft.py 1024 path_for_csv
-srun python benchmarks/mpi4jafft.py 1024 path_for_csv
+srun python benchmarks/pfft3d.py 1024 2_2 NCCL path_for_csv_folder
+srun python benchmarks/jaxfft.py 1024 path_for_csv_folder
+srun python benchmarks/mpi4jafft.py 1024 path_for_csv_folder
 ```
diff --git a/benchmarks/scripts/tablet-script.sh b/benchmarks/scripts/tablet-script.sh
diff --git a/benchmarks/scripts/the_scripter.py b/benchmarks/scripts/the_scripter.py
diff --git a/slurms/v100-32G.slurm → benchmarks/scripts/v100-32G.slurm b/slurms/v100-32G.slurm → benchmarks/scripts/v100-32G.slurm
@@ -51,6 +51,31 @@ export TMPDIR=$JOBSCRATCH
 ln -s $JOBSCRATCH /tmp/nvidia
 
 #
+#!/bin/bash
+
+export OUTPUT_FOLDER_ARGS=0
+
+function run_python() {
+    if [ $# -lt 1 ]; then
+        echo "Usage: run_python <python_script> [arguments for the script]"
+        return 1
+    fi
+
+
+    local script_name=$(basename "$1" .py)
+    local output_dir="output/$script_name"
+
+    if [ $OUTPUT_FOLDER_ARGS -eq 1 ]; then
+        local args=$(echo "${@:2}" | tr ' ' '_')
+        output_dir="output/$script_name/$args"
+    fi
+
+    mkdir -p "$output_dir"
+
+    srun python "$@" > "$output_dir/$script_name.out" 2> "$output_dir/$script_name.err" || true
+}
+
+# To Profile the script 
 function profile_python() {
     if [ $# -lt 1 ]; then
         echo "Usage: profile_python <python_script> [arguments for the script]"
@@ -60,27 +85,29 @@ function profile_python() {
     local script_name=$(basename "$1" .py)
     local output_dir="traces/$script_name"
 
+    if [ $OUTPUT_FOLDER_ARGS -eq 1 ]; then
+        local args=$(echo "${@:2}" | tr ' ' '_')
+        output_dir="traces/$script_name/$args"
+    fi
+
     mkdir -p "$output_dir"
 
     srun nsys profile -t cuda,nvtx,osrt,mpi -o "$output_dir/report_rank%q{SLURM_PROCID}" python "$@" > "$output_dir/$script_name.out" 2> "$output_dir/$script_name.err" || true
 }
 
-
-# Echo of launched commands
-mkdir -p traces/pfft3d
-mkdir -p traces/jaxfft
-
-mkdir -p traces/nsys
-alias srun='srun nsys profile -t cuda,nvtx,osrt,mpi'
+alias slaunch='profile_python'
 
 set -x
 # For the "gpu_p5" partition, the code must be compiled with the compatible modules.
 # Code execution with binding via bind_gpu.sh : 1 GPU per task
 
 # NCCL
-srun nsys profile -t cuda,nvtx,osrt,mpi -o "traces/pfft3d/report_rank%q{SLURM_PROCID}" python benchmarks/pfft3d.py 2_2 1024 NCCL traces/pfft3d > traces/pfft3d/pfft3d.out 2> traces/pfft3d/pfft3d.err || true
-<F2>
-
-srun nsys profile -t cuda,nvtx,osrt,mpi -o "traces/pfft3d/report_rank%q{SLURM_PROCID}" python benchmarks/pfft3d.py 2_2 1024 MPI traces/pfft3d > traces/pfft3d/pfft3d.out 2> traces/pfft3d/pfft3d.err || true
-
-srun nsys profile -t cuda,nvtx,osrt,mpi -o "traces/j/report_rank%q{SLURM_PROCID}" python benchmarks/jaxfft.py 1024 traces/jaxfft > traces/jaxfft/jaxfft.out 2> traces/jaxfft/jaxfft.err || true
+slaunch python benchmarks/pfft3d.py 2_2 1024 NCCL traces/pfft3d
+slaunch python benchmarks/pfft3d.py 4_1 1024 NCCL traces/pfft3d
+slaunch python benchmarks/pfft3d.py 1_4 1024 NCCL traces/pfft3d
+# MPI
+slaunch python benchmarks/pfft3d.py 2_2 1024 MPI traces/pfft3d
+slaunch python benchmarks/pfft3d.py 4_1 1024 MPI traces/pfft3d
+slaunch python benchmarks/pfft3d.py 1_4 1024 MPI traces/pfft3d
+# JAXFFT
+slaunch python benchmarks/jaxfft.py 1024 traces/jaxfft