diff --git a/tests/run_tests.sh b/tests/run_tests.sh new file mode 100755 index 0000000..c651e33 --- /dev/null +++ b/tests/run_tests.sh @@ -0,0 +1,76 @@ +#!/bin/bash + +# Set default parameters +default_grid_size_lat=1 +default_grid_size_lon=1 +default_run_distributed=false + +# Parse command line arguments +while [[ $# -gt 0 ]]; do + case $1 in + -h|--help) + + bold=$(tput bold) + normal=$(tput sgr0) + + echo "Runs the torch-harmonics test suite." + echo "${bold}Arguments:${normal}" + echo " ${bold}-h | --help:${normal} Prints this text." + echo " ${bold}-d | --run_distributed:${normal} Run the distributed test suite." + echo " ${bold}-lat | --grid_size_lat:${normal} Number of ranks in latitudinal direction for distributed case." + echo " ${bold}-lon | --grid_size_lon:${normal} Number of ranks in longitudinal direction for distributed case." + + shift + exit 0 + ;; + -lat|--grid_size_lat) + grid_size_lat="$2" + shift 2 + ;; + -lon|--grid_size_lon) + grid_size_lon="$2" + shift 2 + ;; + -d|--run_distributed) + run_distributed=true + shift + ;; + *) + echo "Unknown argument: $1" + exit 1 + ;; + esac +done + +# Use default values if arguments were not provided +grid_size_lat=${grid_size_lat:-$default_grid_size_lat} +grid_size_lon=${grid_size_lon:-$default_grid_size_lon} +run_distributed=${run_distributed:-$default_run_distributed} + +echo "Running sequential tests:" +python3 -m pytest tests/test_convolution.py tests/test_sht.py + +# Run distributed tests if requested +if [ "$run_distributed" = "true" ]; then + + echo "Running distributed tests with the following parameters:" + echo "Grid size latitude: $grid_size_lat" + echo "Grid size longitude: $grid_size_lon" + + ngpu=$(( ${grid_size_lat} * ${grid_size_lon} )) + + mpirun --allow-run-as-root -np ${ngpu} bash -c " + export CUDA_LAUNCH_BLOCKING=1; + export WORLD_RANK=\${OMPI_COMM_WORLD_RANK}; + export WORLD_SIZE=\${OMPI_COMM_WORLD_SIZE}; + export RANK=\${OMPI_COMM_WORLD_RANK}; + export MASTER_ADDR=localhost; + export MASTER_PORT=29501; + export GRID_H=${grid_size_lat}; + export GRID_W=${grid_size_lon}; + python3 -m pytest tests/test_distributed_sht.py + python3 -m pytest tests/test_distributed_convolution.py + " +else + echo "Skipping distributed tests." +fi