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

@@ -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:

.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:

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

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)
 	}