rocm_jax/build/rocm/run_multi_gpu.sh

#!/usr/bin/env bash
# Copyright 2022 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -xu

# Function to run tests with specified GPUs
run_tests() {
    local base_dir=./logs
    local gpu_devices="$1"
    export HIP_VISIBLE_DEVICES=$gpu_devices
    python3 -m pytest --html=$base_dir/multi_gpu_pmap_test_log.html --reruns 3  tests/pmap_test.py
    python3 -m pytest --html=$base_dir/multi_gpu_multi_device_test_log.html --reruns 3 tests/multi_device_test.py
    python3 -m pytest_html_merger -i $base_dir/ -o  $base_dir/final_compiled_report.html
}

# Check for required commands
if ! command -v lspci &> /dev/null; then
    echo "lspci command not found, aborting."
    exit 1
fi

if ! command -v python3 &> /dev/null; then
    echo "Python3 is not available, aborting."
    exit 1
fi

# GPU detection and test execution
gpu_count=$(lspci | grep -c 'controller.*AMD/ATI')
echo "Number of AMD/ATI GPUs detected: $gpu_count"

if [[ $gpu_count -gt 8 ]]; then
    run_tests "0,1,2,3,4,5,6,7"
elif [[ $gpu_count -gt 4 ]]; then
    run_tests "0,1,2,3"
elif [[ $gpu_count -gt 2 ]]; then
    run_tests "0,1"
else
    run_tests "0"
fi