From 5fdedb52765715be6be8a594bbe5a68cbf529ece Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Mon, 16 Dec 2024 17:21:01 -0800
Subject: [PATCH] Parallelize BuildKite "All GPU drivers test".

Each shard takes on a subset of the supported driver versions, as determined
by a counter.

Also sort the order in which the list of supported versions is written out
so that console output isn't confusing.

PiperOrigin-RevId: 706887043
---
 .buildkite/pipeline.yaml            |  1 +
 .buildkite/release.yaml             |  1 +
 tools/gpu/all_drivers_test.sh       | 40 ++++++++++++++++++++++++-----
 tools/gpu/drivers/install_driver.go |  2 ++
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/.buildkite/pipeline.yaml b/.buildkite/pipeline.yaml
index a708b9978..8d6218950 100644
--- a/.buildkite/pipeline.yaml
+++ b/.buildkite/pipeline.yaml
@@ -213,6 +213,7 @@ steps:
   - <<: *common
     <<: *source_test_continuous
     label: ":screwdriver: All GPU Drivers Test"
+    parallelism: 8
     commands:
       - tools/gpu/all_drivers_test.sh
     agents:
diff --git a/.buildkite/release.yaml b/.buildkite/release.yaml
index ef3e1ccce..53f4a6276 100644
--- a/.buildkite/release.yaml
+++ b/.buildkite/release.yaml
@@ -77,6 +77,7 @@ steps:
       queue: gpu
   - <<: *common
     label: ":screwdriver: All GPU Drivers Test"
+    parallelism: 8
     commands:
       - tools/gpu/all_drivers_test.sh
     agents:
diff --git a/tools/gpu/all_drivers_test.sh b/tools/gpu/all_drivers_test.sh
index ed1f7aa1c..2f3bf0fe9 100755
--- a/tools/gpu/all_drivers_test.sh
+++ b/tools/gpu/all_drivers_test.sh
@@ -14,28 +14,56 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Script to easily run gpu tests on all supported driver versions. This should
+# Script to easily run GPU tests on all supported driver versions. This should
 # be run from the gVisor repo root directory.
 set -ueo pipefail
 
-tmp_file=$(mktemp)
+tmp_file="$(mktemp)"
 trap "rm -f ${tmp_file}" EXIT
 
 make sudo TARGETS=tools/gpu:main ARGS="list --outfile=${tmp_file}"
-read -r -a versions <<< "$(cat "${tmp_file}")"
+read -r -a all_versions <<< "$(cat "${tmp_file}")"
+
+if [[ "${#all_versions[@]}" -eq 0 ]]; then
+  echo 'No driver versions found.' >&2
+  exit 1
+fi
+
+# https://buildkite.com/docs/pipelines/tutorials/parallel-builds
+my_shard="${BUILDKITE_PARALLEL_JOB:-0}"
+total_shards="${BUILDKITE_PARALLEL_JOB_COUNT:-1}"
+
+counter=0
+versions=()
+for driver in "${all_versions[@]}"; do
+  modulo="$(( "$counter" % "$total_shards" ))"
+  explanation="${counter} % ${total_shards} == ${modulo}; we are shard ${my_shard} of $(( ${total_shards} - 1 ))"
+  if [[ "$modulo" -eq "$my_shard" ]]; then
+    echo "Will test driver ${driver} ($explanation)" >&2
+    versions+=("$driver")
+  else
+    echo "Skipping driver ${driver} ($explanation)" >&2
+  fi
+  counter="$(( "$counter" + 1 ))"
+done
+
+if [[ "${#versions[@]}" -eq 0 ]]; then
+  echo "No versions to test on this shard (we are shard ${my_shard} of $(( ${total_shards} - 1 )))." >&2
+  exit 0
+fi
 
 num_successful=0
 for driver in "${versions[@]}"; do
   set +e
   make sudo TARGETS=tools/gpu:main ARGS="install --version ${driver}"
-  install_exit_code=$?
+  install_exit_code="$?"
   set -e
-  if [[ $install_exit_code -ne 0 ]]; then
+  if [[ "$install_exit_code" -ne 0 ]]; then
     echo "Installing driver ${driver} failed. Not testing this version." >&2
     continue
   fi
   make gpu-smoke-tests
-  num_successful="$(( $num_successful + 1 ))"
+  num_successful="$(( "$num_successful" + 1 ))"
 done
 if [[ "$num_successful" == 0 ]]; then
   echo 'No version was successfully tested.' >&2
diff --git a/tools/gpu/drivers/install_driver.go b/tools/gpu/drivers/install_driver.go
index 193644ce6..091f29401 100644
--- a/tools/gpu/drivers/install_driver.go
+++ b/tools/gpu/drivers/install_driver.go
@@ -23,6 +23,7 @@ import (
 	"net/http"
 	"os"
 	"os/exec"
+	"sort"
 	"strings"
 
 	"gvisor.dev/gvisor/pkg/log"
@@ -219,6 +220,7 @@ func ListSupportedDrivers(outfile string) error {
 	nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, checksum string) {
 		list = append(list, version.String())
 	})
+	sort.Strings(list)
 	if _, err := out.WriteString(strings.Join(list, " ") + "\n"); err != nil {
 		return fmt.Errorf("failed to write to outfile: %w", err)
 	}