Added script for scaling tests

Author: Thomas

pySDC/projects/GPU/analysis_scripts/parallel_scaling.py

@@ -0,0 +1,170 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pickle
+from pySDC.helpers.stats_helper import get_sorted
+from pySDC.projects.GPU.configs.base_config import get_config
+from pySDC.projects.GPU.etc.generate_jobscript import write_jobscript, PROJECT_PATH
+
+
+class ScalingConfig(object):
+    cluster = None
+    config = ''
+    base_resolution = -1
+    base_resolution_weak = -1
+    useGPU = False
+    partition = None
+    tasks_per_node = None
+    ndim = 2
+    tasks_time = 1
+    max_steps_space = None
+    max_steps_space_weak = None
+
+    def __init__(self, space_time_parallel):
+        if space_time_parallel in ['False', False]:
+            self._tasks_time = 1
+        else:
+            self._tasks_time = self.tasks_time
+
+    def get_resolution_and_tasks(self, strong, i):
+        if strong:
+            return self.base_resolution, [1, self._tasks_time, 2**i]
+        else:
+            return self.base_resolution_weak * (2**i), [1, self._tasks_time, (2 * self.ndim) ** i]
+
+    def run_scaling_test(self, strong=True):
+        max_steps = self.max_steps_space if strong else self.max_steps_space_weak
+        for i in range(max_steps):
+            res, procs = self.get_resolution_and_tasks(strong, i)
+
+            sbatch_options = [f'-n {np.prod(procs)}', f'-p {self.partition}']
+            if self.useGPU:
+                srun_options = ['--cpus-per-task=4', '--gpus-per-task=1']
+                sbatch_options += ['--cpus-per-task=4', '--gpus-per-task=1']
+            else:
+                srun_options = []
+
+            procs = (''.join(f'{me}/' for me in procs))[:-1]
+            command = f'run_experiment.py --mode=run --res={res} --config={self.config} --procs={procs}'
+
+            if self.useGPU:
+                command += ' --useGPU=True'
+
+            write_jobscript(sbatch_options, srun_options, command, self.cluster)
+
+    def plot_scaling_test(self, strong, ax, plot_ideal=False, **plotting_params):
+        timings = {}
+
+        max_steps = self.max_steps_space if strong else self.max_steps_space_weak
+        for i in range(max_steps):
+            res, procs = self.get_resolution_and_tasks(strong, i)
+
+            args = {'useGPU': self.useGPU, 'config': self.config, 'res': res, 'procs': procs, 'mode': None}
+
+            config = get_config(args)
+
+            path = f'data/{config.get_path(ranks=[me -1 for me in procs])}-stats-whole-run.pickle'
+            with open(path, 'rb') as file:
+                stats = pickle.load(file)
+
+            timing_step = get_sorted(stats, type='timing_step')
+
+            timings[np.prod(procs) / self.tasks_per_node] = np.mean([me[1] for me in timing_step])
+
+        ax.loglog(timings.keys(), timings.values(), **plotting_params)
+        if plot_ideal:
+            ax.loglog(
+                timings.keys(),
+                list(timings.values())[0] * list(timings.keys())[0] / np.array(list(timings.keys())),
+                ls='--',
+                color='grey',
+                label='ideal',
+            )
+        ax.set_xlabel(r'$N_\mathrm{nodes}$')
+        ax.set_ylabel(r'$t_\mathrm{step}$')
+
+
+class CPUConfig(ScalingConfig):
+    cluster = 'jusuf'
+    partition = 'batch'
+    tasks_per_node = 128
+
+
+class GPUConfig(ScalingConfig):
+    cluster = 'booster'
+    partition = 'booster'
+    tasks_per_node = 4
+    useGPU = True
+
+
+class GrayScottSpaceScalingCPU(CPUConfig, ScalingConfig):
+    base_resolution = 2048
+    base_resolution_weak = 256
+    config = 'GS_scaling'
+    max_steps_space = 10
+    max_steps_space_weak = 6
+    tasks_time = 3
+
+
+class GrayScottSpaceScalingGPU(GPUConfig, ScalingConfig):
+    base_resolution_weak = 256 * 32
+    base_resolution = 2048
+    config = 'GS_scaling'
+    max_steps_space = 4
+    max_steps_space_weak = 4
+    tasks_time = 3
+
+
+def plot_scalings(strong, problem, kwargs):
+    if problem == 'GS':
+        fig, ax = plt.subplots()
+
+        plottings_params = [
+            {'plot_ideal': strong, 'marker': 'x', 'label': 'CPU'},
+            {'marker': '>', 'label': 'CPU space time parallel'},
+            {'marker': '^', 'label': 'GPU'},
+        ]
+        configs = [
+            GrayScottSpaceScalingCPU(space_time_parallel=False),
+            GrayScottSpaceScalingCPU(space_time_parallel=True),
+            GrayScottSpaceScalingGPU(space_time_parallel=False),
+        ]
+
+        for config, params in zip(configs, plottings_params):
+            config.plot_scaling_test(strong=strong, ax=ax, **params)
+        ax.legend(frameon=False)
+        fig.savefig(f'{PROJECT_PATH}/plots/{"strong" if strong else "weak"}_scaling_{problem}.pdf', bbox_inches='tight')
+    else:
+        raise NotImplementedError
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--scaling', type=str, choices=['strong', 'weak'], default='strong')
+    parser.add_argument('--mode', type=str, choices=['run', 'plot'], default='run')
+    parser.add_argument('--problem', type=str, default='GS')
+    parser.add_argument('--XPU', type=str, choices=['CPU', 'GPU'], default='CPU')
+    parser.add_argument('--space_time', type=str, choices=['True', 'False'], default='False')
+
+    args = parser.parse_args()
+
+    strong = args.scaling == 'strong'
+
+    if args.problem == 'GS':
+        if args.XPU == 'CPU':
+            configClass = GrayScottSpaceScalingCPU
+        else:
+            configClass = GrayScottSpaceScalingGPU
+    else:
+        raise NotImplementedError(f'Don\'t know problem {args.problem!r}')
+
+    kwargs = {'space_time_parallel': args.space_time}
+    config = configClass(**kwargs)
+
+    if args.mode == 'run':
+        config.run_scaling_test(strong=strong)
+    elif args.mode == 'plot':
+        plot_scalings(strong=strong, problem=args.problem, kwargs=kwargs)
+    else:
+        raise NotImplementedError(f'Don\'t know mode {args.mode!r}')

pySDC/projects/GPU/configs/GS_configs.py

@@ -12,6 +12,8 @@ def get_config(args):
         return GrayScott_GoL(args)
     elif name == 'GS_USkate':
         return GrayScott_USkate(args)
+    elif name == 'GS_scaling':
+        return GrayScottScaling(args)
     else:
         return NotImplementedError(f'Don\'t know config {name}')
 
@@ -183,3 +185,18 @@ def get_description(self, *args, **kwargs):
         desc['convergence_controllers'][Adaptivity] = {'e_tol': 1e-3}
         self.Tend = 100000
         return desc
+
+
+class GrayScottScaling(GrayScott):
+    def get_description(self, *args, **kwargs):
+        desc = super().get_description(*args, **kwargs)
+        desc['problem_params']['L'] = 2
+        desc['problem_params']['num_blobs'] = 4
+        desc['sweeper_params']['skip_residual_computation'] = ('IT_CHECK', 'IT_DOWN', 'IT_UP', 'IT_FINE', 'IT_COARSE')
+        self.Tend = 100 * desc['level_params']['dt']
+        return desc
+
+    def get_controller_params(self, *args, **kwargs):
+        params = super().get_controller_params(*args, **kwargs)
+        params['hook_class'] = []
+        return params

pySDC/projects/GPU/etc/generate_jobscript.py

@@ -0,0 +1,48 @@
+PROJECT_PATH = '/p/project1/ccstma/baumann7/pySDC/pySDC/projects/GPU'
+DEFAULT_SBATCH_OPTIONS = ['-A cstma', '--threads-per-core=1', f'--output={PROJECT_PATH}/etc/slurm-out/%j.txt']
+DEFAULT_SRUN_OPTIONS = ['--cpu-bind=sockets']
+
+
+def generate_directories():
+    import os
+
+    for name in ['jobscripts', 'slurm-out']:
+        path = f'{PROJECT_PATH}/etc/{name}'
+        os.makedirs(path, exist_ok=True)
+
+
+def get_jobscript_text(sbatch_options, srun_options, command, cluster):
+    msg = '#!/usr/bin/bash\n\n'
+    for op in DEFAULT_SBATCH_OPTIONS + sbatch_options:
+        msg += f'#SBATCH {op}\n'
+
+    msg += f'\nsource {PROJECT_PATH}/etc/venv_{cluster.lower()}/activate.sh\n'
+
+    srun_cmd = 'srun'
+    for op in DEFAULT_SRUN_OPTIONS + srun_options:
+        srun_cmd += f' {op}'
+
+    msg += f'\n{srun_cmd} python {PROJECT_PATH}/{command}'
+    return msg
+
+
+def write_jobscript(sbatch_options, srun_options, command, cluster, submit=True):
+    generate_directories()
+
+    text = get_jobscript_text(sbatch_options, srun_options, command, cluster)
+
+    path = f'{PROJECT_PATH}/etc/jobscripts/{command.replace(" ", "").replace("/", "_")}-{cluster}.sh'
+    with open(path, 'w') as file:
+        file.write(text)
+
+    if submit:
+        import os
+
+        os.system(f'sbatch {path}')
+
+
+if __name__ == '__main__':
+    sbatch_options = ['--nodes=1']
+    srun_options = []
+    command = 'run_problems.py'
+    write_jobscript(sbatch_options, srun_options, command, 'jusuf')

pySDC/projects/GPU/etc/venv_booster/activate.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+# See https://stackoverflow.com/a/28336473
+SOURCE_PATH="${BASH_SOURCE[0]:-${(%):-%x}}"
+
+RELATIVE_PATH="$(dirname "$SOURCE_PATH")"
+ABSOLUTE_PATH="$(realpath "${RELATIVE_PATH}")"
+
+[[ "$0" != "${SOURCE_PATH}" ]] && echo "The activation script must be sourced, otherwise the virtual environment will not work." || ( echo "Vars script must be sourced." && exit 1) ;
+
+source "${ABSOLUTE_PATH}"/config.sh
+source "${ABSOLUTE_PATH}"/modules.sh
+
+export PYTHONPATH="$(echo "${ENV_DIR}"/lib/python*/site-packages):${PYTHONPATH}"
+
+source "${ENV_DIR}"/bin/activate

pySDC/projects/GPU/etc/venv_booster/config.sh

@@ -0,0 +1,12 @@
+SOURCE_PATH="${BASH_SOURCE[0]:-${(%):-%x}}"
+
+## Check if this script is sourced
+[[ "$0" != "${SOURCE_PATH}" ]] && echo "Setting vars" || ( echo "Vars script must be sourced." && exit 1) ;
+## Determine location of this file
+RELATIVE_PATH="$(dirname "$SOURCE_PATH")"
+ABSOLUTE_PATH="$(realpath "${RELATIVE_PATH}")"
+####################################
+
+### User Configuration
+export ENV_NAME="GPU"             # Default Name of the venv is the directory that contains this file
+export ENV_DIR="${ABSOLUTE_PATH}"/venv         # Default location of this VENV is "./venv"

pySDC/projects/GPU/etc/venv_booster/create_kernel.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+
+SOURCE_PATH="${BASH_SOURCE[0]:-${(%):-%x}}"
+
+RELATIVE_PATH="$(dirname "$SOURCE_PATH")"
+ABSOLUTE_PATH="$(realpath "${RELATIVE_PATH}")"
+source "${ABSOLUTE_PATH}"/config.sh
+
+KERNELFILE="${ENV_DIR}"/kernel.sh
+
+echo the name is "$ENV_NAME"
+
+echo "Setting up the kernel script in the following dir: " "${KERNELFILE}"
+echo "If you use multiprocessing, edit this file to remove the srun call from the kernel and run it again."
+echo '#!/bin/bash
+
+source "'"${ABSOLUTE_PATH}"'"/activate.sh
+
+hostname=$(hostname)
+
+if [[ $hostname == *"login"* || $hostname == *"jsfl"* ]]; then
+   exec python -m ipykernel "$@"
+else
+   srun python -m ipykernel "$@"
+fi
+' > "${KERNELFILE}"
+
+chmod a+x "${KERNELFILE}"
+
+mkdir -p ~/.local/share/jupyter/kernels/"${ENV_NAME}"
+echo '{
+ "argv": [
+  "'"${KERNELFILE}"'",
+  "-f",
+  "{connection_file}"
+ ],
+ "display_name": "'"${ENV_NAME}"'",
+ "language": "python"
+}' > ~/.local/share/jupyter/kernels/"${ENV_NAME}"/kernel.json

pySDC/projects/GPU/etc/venv_booster/create_python_for_vscode.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+SOURCE_PATH="${BASH_SOURCE[0]:-${(%):-%x}}"
+
+RELATIVE_PATH="$(dirname "$SOURCE_PATH")"
+ABSOLUTE_PATH="$(realpath "${RELATIVE_PATH}")"
+source "${ABSOLUTE_PATH}"/config.sh
+PYTHONWRAPPER="${ABSOLUTE_PATH}"/python
+
+echo '#!/bin/bash
+module purge 2> /dev/null
+deactivate 2> /dev/null
+source '"'${ABSOLUTE_PATH}'"'/activate.sh
+python "$@"
+' > "${PYTHONWRAPPER}"
+
+chmod a+x "${PYTHONWRAPPER}"

pySDC/projects/GPU/etc/venv_booster/modules.sh

@@ -0,0 +1,13 @@
+module --force purge
+module load Stages/2024
+module load GCC
+module load ParaStationMPI
+module load NCCL
+module load MPI-settings/CUDA
+module load UCX-settings/RC-CUDA
+module load Python
+module load CuPy
+module load FFTW
+module load mpi4py
+module load FFmpeg/.6.0
+module load SciPy-Stack