From abc46770a98777a9fd1b072adcf8becb76bfe125 Mon Sep 17 00:00:00 2001
From: Stephen Sachs <ssachs@nvidia.com>
Date: Fri, 26 Sep 2025 14:34:23 +0200
Subject: [PATCH] Check if sufficient GPUs are available

The CUDA error message "Test CUDA failure util.cu:706 'invalid device ordinal'"
is not as helpful. Test this explicitly and guide the user.
---
 src/common.cu | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/common.cu b/src/common.cu
index 8dc16cd..4b8e8c0 100644
--- a/src/common.cu
+++ b/src/common.cu
@@ -1137,9 +1137,21 @@ testResult_t run() {
   size_t maxMem = ~0;
   char* envstr = getenv("NCCL_TESTS_DEVICE");
   int gpu0 = envstr ? atoi(envstr) : -1;
+  int available_devices;
+  CUDACHECK(cudaGetDeviceCount(&available_devices));
   for (int i=0; i<nThreads*nGpus; i++) {
     int cudaDev = (gpu0 != -1 ? gpu0 : localRank*nThreads*nGpus) + i;
     int rank = proc*nThreads*nGpus+i;
+    if (cudaDev >= available_devices) {
+      fprintf(stderr,
+              "Invalid number of GPUs: %d requested but only %d were found.\n",
+              (gpu0 != -1 ? gpu0 : localRank * nThreads * nGpus) +
+                  nThreads * nGpus,
+              available_devices);
+      fprintf(stderr,
+              "Please check the number of processes and GPUs per process.\n");
+      return testNotImplemented;
+    }
     cudaDeviceProp prop;
     CUDACHECK(cudaGetDeviceProperties(&prop, cudaDev));
     len += snprintf(line+len, MAX_LINE-len, "#  Rank %2d Group %2d Pid %6d on %10s device %2d [%04x:%02x:%02x] %s\n",