Extend NCCL test to p5.48xlarge

Extend the possibility to run NCCL test with p5.48xlarge. The case has not been added to the integration tests config.

Signed-off-by: Luca Carrogu <carrogu@amazon.com>

Author: lukeseawalker

tests/integration-tests/tests/efa/test_efa.py

@@ -58,8 +58,8 @@ def test_efa(
 
     _test_shm_transfer_is_enabled(scheduler_commands, remote_command_executor, partition="efa-enabled")
 
-    if instance == "p4d.24xlarge" and os != "centos7":
-        _test_nccl_benchmarks(remote_command_executor, test_datadir, "openmpi", scheduler_commands)
+    if instance in ["p4d.24xlarge", "p5.48xlarge"] and os != "centos7":
+        _test_nccl_benchmarks(remote_command_executor, test_datadir, "openmpi", scheduler_commands, instance)
 
     assert_no_errors_in_logs(remote_command_executor, scheduler, skip_ice=True)
 
@@ -102,7 +102,7 @@ def _test_shm_transfer_is_enabled(scheduler_commands, remote_command_executor, p
     assert_that(result.stdout).does_not_contain("SHM transfer will be disabled because of ptrace protection")
 
 
-def _test_nccl_benchmarks(remote_command_executor, test_datadir, mpi_module, scheduler_commands):
+def _test_nccl_benchmarks(remote_command_executor, test_datadir, mpi_module, scheduler_commands, instance):
     logging.info("Running NCCL benchmarks")
     remote_command_executor.run_remote_script(
         str(test_datadir / "nccl_benchmarks" / "init_nccl_benchmarks.sh"), args=[mpi_module], hide=True, timeout=600
@@ -139,5 +139,15 @@ def _test_nccl_benchmarks(remote_command_executor, test_datadir, mpi_module, sch
         "cat /shared/nccl_tests.out | grep -E '1073741824\\s+268435456' | awk '{print $12}'"
     ).stdout
 
-    # Expected "in-place busbw" bandwidth with 2 nodes, 8 tasks per node is about 27GB/s
-    assert_that(float(max_bandwidth)).is_greater_than(26.0)
+    instance_bandwidth_dict = {
+        # p4d.24xlarge - Expected "in-place busbw" bandwidth with 2 nodes, 8 tasks per node is about 27GB/s
+        "p4d.24xlarge": 26.0,
+        # p5.48xlarge - Expected "in-place busbw" bandwidth with 2 nodes, 8 tasks per node is about 250GB/s
+        "p5.48xlarge": 250.0,
+    }
+
+    expected_bandwidth = instance_bandwidth_dict.get(instance)
+    if expected_bandwidth is None:
+        pytest.fail(f"Instance {instance} is not valid for multiple bandwidth tests")
+
+    assert_that(float(max_bandwidth)).is_greater_than(expected_bandwidth)

tests/integration-tests/tests/efa/test_efa/test_efa/nccl_benchmarks/init_nccl_benchmarks.sh

@@ -1,5 +1,6 @@
 #!/usr/bin/env bash
-set -e
+exec &>/shared/init_nccl.out
+set -xe
 
 rm -rf /shared/${1}
 
@@ -8,7 +9,7 @@ NCCL_BENCHMARKS_VERSION='2.13.8'
 NCCL_VERSION='2.19.4-1'
 OFI_NCCL_VERSION='1.7.4-aws'
 MPI_HOME=$(which mpirun | awk -F '/bin' '{print $1}')
-NVCC_GENCODE="-gencode=arch=compute_80,code=sm_80" # Arch for NVIDIA A100
+NVCC_GENCODE="-gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_90,code=sm_90 -gencode=arch=compute_90,code=compute_90" # Arch for NVIDIA A100 and H100, ref https://docs.nvidia.com/cuda/ada-compatibility-guide/index.html
 
 mkdir -p /shared/${1}
 
@@ -24,7 +25,7 @@ cd /shared/${1}
 wget https://github.com/NVIDIA/nccl-tests/archive/v${NCCL_BENCHMARKS_VERSION}.tar.gz
 tar zxvf "v${NCCL_BENCHMARKS_VERSION}.tar.gz"
 cd "nccl-tests-${NCCL_BENCHMARKS_VERSION}/"
-NVCC_GENCODE="${NVCC_GENCODE}" make MPI=1 MPI_HOME=${MPI_HOME} NCCL_HOME=/shared/${1}/nccl-${NCCL_VERSION}/build/
+NVCC_GENCODE="${NVCC_GENCODE}" make MPI=1 MPI_HOME=${MPI_HOME} NCCL_HOME=/shared/${1}/nccl-${NCCL_VERSION}/build/ CUDA_HOME=/usr/local/cuda
 
 wget https://github.com/aws/aws-ofi-nccl/archive/v${OFI_NCCL_VERSION}.tar.gz
 tar xvfz v${OFI_NCCL_VERSION}.tar.gz