[CI] Use AWS for GPU instances

.buildkite/pipelines/cuvs-snapshot/run-tests.yml

@@ -3,11 +3,12 @@ steps:
     command: .buildkite/scripts/cuvs-snapshot/run-gradle.sh -Dbwc.checkout.align=true -Dorg.elasticsearch.build.cache.push=true -Dignore.tests.seed -Dscan.capture-file-fingerprints {{matrix.GRADLE_TASK}}
     timeout_in_minutes: 300
     agents:
-      provider: gcp
-      image: family/elasticsearch-ubuntu-2404-nvidia
-      machineType: g2-standard-32
-      buildDirectory: /dev/shm/bk
-      zones: us-central1-b,us-central1-c
+      provider: aws
+      imagePrefix: elasticsearch-aws-ubuntu-2204-nvidia
+      instanceType: g6.8xlarge
+      diskSizeGb: 350
+      diskType: gp3
+      diskName: /dev/sda1
     env:
       GRADLE_TASK: "{{matrix.GRADLE_TASK}}"
     matrix:

.buildkite/pipelines/cuvs-snapshot/update-snapshot.yml

@@ -2,11 +2,12 @@ steps:
   - label: "Smoke test and update new cuVS snapshot"
     command: .buildkite/scripts/cuvs-snapshot/update-current-snapshot-version.sh
     agents:
-      provider: gcp
-      image: family/elasticsearch-ubuntu-2404-nvidia
-      machineType: g2-standard-16
-      zones: us-central1-b,us-central1-c
-      diskSizeGb: 150
+      provider: aws
+      imagePrefix: elasticsearch-aws-ubuntu-2204-nvidia
+      instanceType: g6.2xlarge
+      diskSizeGb: 350
+      diskType: gp3
+      diskName: /dev/sda1
   - wait: ~
   - trigger: "elasticsearch-cuvs-run-tests"
     build:

.buildkite/pipelines/pull-request/gpu.yml

@@ -10,11 +10,12 @@ steps:
         command: .buildkite/scripts/cuvs-snapshot/run-gradle.sh -Dbwc.checkout.align=true -Dorg.elasticsearch.build.cache.push=true -Dignore.tests.seed -Dscan.capture-file-fingerprints {{matrix.GRADLE_TASK}}
         timeout_in_minutes: 300
         agents:
-          provider: gcp
-          image: family/elasticsearch-ubuntu-2404-nvidia
-          machineType: g2-standard-32
-          buildDirectory: /dev/shm/bk
-          zones: us-central1-b,us-central1-c
+          provider: aws
+          imagePrefix: elasticsearch-aws-ubuntu-2204-nvidia
+          instanceType: g6.8xlarge
+          diskSizeGb: 350
+          diskType: gp3
+          diskName: /dev/sda1
         env:
           GRADLE_TASK: "{{matrix.GRADLE_TASK}}"
         matrix:

.buildkite/scripts/cuvs-snapshot/configure.sh

@@ -1,7 +1,5 @@
 #!/bin/bash
 
-set -euo pipefail
-
 if [[ -f /etc/profile.d/elastic-nvidia.sh ]]; then
   export JAVA_HOME="$HOME/.java/openjdk24"
   export PATH="$JAVA_HOME/bin:$PATH"
@@ -10,34 +8,35 @@ if [[ -f /etc/profile.d/elastic-nvidia.sh ]]; then
 
   export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}"
   source /etc/profile.d/elastic-nvidia.sh
-fi
-
-# Not running this before the tests results in an error when running the tests
-# No idea why...
-nvidia-smi
 
-CURRENT_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-ELASTICSEARCH_REPO_DIR="$(cd "$CURRENT_SCRIPT_DIR/../../.." && pwd)"
-
-CUVS_SNAPSHOT_VERSION="${CUVS_SNAPSHOT_VERSION:-$(cat "$CURRENT_SCRIPT_DIR"/current-snapshot-version)}"
-CUVS_ARCHIVE="cuvs-$CUVS_SNAPSHOT_VERSION.tar.gz"
-CUVS_URL="https://storage.googleapis.com/elasticsearch-cuvs-snapshots/$CUVS_ARCHIVE"
+  # Not running this before the tests results in an error when running the tests
+  # No idea why...
+  if [[ "${BUILDKITE:-}" != "" && "${CI:-}" == "true" ]]; then
+    nvidia-smi
+  fi
+fi
 
-CUVS_WORKSPACE=${CUVS_WORKSPACE:-$(cd "$(mktemp -d)")}
-CUVS_DIR="$(pwd)/cuvs-$CUVS_SNAPSHOT_VERSION"
+LIBCUVS_GCS_BUCKET="elasticsearch-cuvs-snapshots"
 
-curl -O "$CUVS_URL"
-tar -xzf "$CUVS_ARCHIVE"
+LIBCUVS_DIR="/opt/libcuvs"
+mkdir -p "$LIBCUVS_DIR"
+chmod 777 "$LIBCUVS_DIR"
 
-CUVS_VERSION=$(cd "$CUVS_DIR/cuvs-java/target" && mvn help:evaluate -Dexpression=project.version -q -DforceStdout)
+CUVS_VERSION=$(grep 'cuvs_java' build-tools-internal/version.properties | awk '{print $3}')
 
-LD_LIBRARY_PATH=$(echo "$LD_LIBRARY_PATH" | tr ':' '\n' | grep -v "libcuvs/linux-x64" | tr '\n' ':' | sed 's/:$//')
-LD_LIBRARY_PATH="$CUVS_DIR/libcuvs/linux-x64:$LD_LIBRARY_PATH"
-export LD_LIBRARY_PATH
+LIBCUVS_VERSION_DIR="$LIBCUVS_DIR/$CUVS_VERSION"
 
-cd "$CUVS_DIR/cuvs-java/target"
-mvn install:install-file -Dfile="cuvs-java-$CUVS_VERSION.jar" -DartifactId=elastic-cuvs-java -DgeneratePom=true
+if [[ ! -d "$LIBCUVS_VERSION_DIR" ]]; then
+  mkdir -p $LIBCUVS_VERSION_DIR
+  cd "$LIBCUVS_VERSION_DIR"
+  CUVS_ARCHIVE="libcuvs-$CUVS_VERSION.tar.gz"
+  curl -fO "https://storage.googleapis.com/$LIBCUVS_GCS_BUCKET/libcuvs/$CUVS_ARCHIVE"
+  tar -xzf "$CUVS_ARCHIVE"
+  rm -f "$CUVS_ARCHIVE"
+  if [[ -d "$CUVS_VERSION" ]]; then
+    mv "$CUVS_VERSION/*" ./
+  fi
+  cd -
+fi
 
-cd "$ELASTICSEARCH_REPO_DIR"
-PLUGIN_GRADLE_FILE=x-pack/plugin/gpu/build.gradle
-sed -i "s|implementation 'com.nvidia.cuvs:elastic-cuvs-java:.*'|implementation 'com.nvidia.cuvs:elastic-cuvs-java:$CUVS_VERSION'|" "$PLUGIN_GRADLE_FILE"
+export LD_LIBRARY_PATH="$LIBCUVS_VERSION_DIR:${LD_LIBRARY_PATH:-}"

.buildkite/scripts/cuvs-snapshot/remove-verification-metadata.py

@@ -0,0 +1,13 @@
+import re
+
+regex = re.compile(
+    r'<component group="com.nvidia.cuvs" name="cuvs-.*?</component>\s*',
+    re.MULTILINE | re.DOTALL,
+)
+
+with open("gradle/verification-metadata.xml", "r+") as f:
+    text = f.read()
+    text = regex.sub("", text)
+    f.seek(0)
+    f.truncate()
+    f.write(text)

.buildkite/scripts/cuvs-snapshot/update-current-snapshot-version.sh

@@ -2,35 +2,35 @@
 
 set -euo pipefail
 
-SNAPSHOT_VERSION_FILE=.buildkite/scripts/cuvs-snapshot/current-snapshot-version
-BRANCH_TO_UPDATE="${BRANCH_TO_UPDATE:-${BUILDKITE_BRANCH:-cuvs-snapshot}}"
+echo "--- Updating cuvs-java version"
 
-if [[ -z "${CUVS_SNAPSHOT_VERSION:-}" ]]; then
-  echo "CUVS_SNAPSHOT_VERSION not set. Set this to update the current snapshot version."
-  exit 1
-fi
+git checkout "$BUILDKITE_BRANCH"
+git pull --ff-only origin "$BUILDKITE_BRANCH"
 
-if [[ "$CUVS_SNAPSHOT_VERSION" == "$(cat $SNAPSHOT_VERSION_FILE)" ]]; then
-  echo "Current snapshot version already set to '$CUVS_SNAPSHOT_VERSION'. No need to update."
-  exit 0
-fi
+# Replace `cuvs_java = <version>` string in version.properties and maintain the same indentation
+sed -E "s/^(cuvs_java *= *[^ ]*  *).*\$/\1$CUVS_JAVA_VERSION/" build-tools-internal/version.properties > new-version.properties
+mv new-version.properties build-tools-internal/version.properties
 
-echo "--- Configuring libcuvs/cuvs-java"
-source .buildkite/scripts/cuvs-snapshot/configure.sh
+python3 .buildkite/scripts/lucene-snapshot/remove-verification-metadata.py
+./gradlew --write-verification-metadata sha256
 
 if [[ "${SKIP_TESTING:-}" != "true" ]]; then
-  echo "--- Testing snapshot before updating"
+  echo "--- Testing cuvs-java before committing"
   ./gradlew -Druntime.java=24 :x-pack:plugin:gpu:yamlRestTest -S
 fi
 
-echo "--- Updating snapshot"
+if git diff-index --quiet HEAD --; then
+  echo 'No changes to commit.'
+  exit 0
+fi
+
+echo "--- Committing changes"
 
-echo "$CUVS_SNAPSHOT_VERSION" > "$SNAPSHOT_VERSION_FILE"
+git config --global user.name elasticsearchmachine
+git config --global user.email '[email protected]'
 
-CURRENT_SHA="$(gh api "/repos/elastic/elasticsearch/contents/$SNAPSHOT_VERSION_FILE?ref=$BRANCH_TO_UPDATE" | jq -r .sha)" || true
+git add build-tools-internal/version.properties
+git add gradle/verification-metadata.xml
 
-gh api -X PUT "/repos/elastic/elasticsearch/contents/$SNAPSHOT_VERSION_FILE" \
-  -f branch="$BRANCH_TO_UPDATE" \
-  -f message="Update cuvs snapshot version to $CUVS_VERSION" \
-  -f content="$(base64 -w 0 "$WORKSPACE/$SNAPSHOT_VERSION_FILE")" \
-  -f sha="$CURRENT_SHA"
+git commit -m "[Automated] Update cuvs-java to $CUVS_JAVA_VERSION"
+git push origin "$BUILDKITE_BRANCH"

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/RepositoriesSetupPlugin.java

@@ -22,6 +22,7 @@
 public class RepositoriesSetupPlugin implements Plugin<Project> {
 
     private static final Pattern LUCENE_SNAPSHOT_REGEX = Pattern.compile("\\w+-snapshot-([a-z0-9]+)");
+    private static final Pattern CUVS_SNAPSHOT_REGEX = Pattern.compile("([0-9\\.]+-[a-z0-9]+)-SNAPSHOT");
 
     @Override
     public void apply(Project project) {
@@ -60,5 +61,17 @@ public static void configureRepositories(Project project) {
                 exclusiveRepo.forRepositories(luceneRepo);
             });
         }
+
+        String cuvsVersion = VersionProperties.getVersions().get("cuvs_java");
+        if (cuvsVersion.contains("-SNAPSHOT")) {
+            MavenArtifactRepository cuvsRepo = repos.maven(repo -> {
+                repo.setName("cuvs-snapshots");
+                repo.setUrl("https://storage.googleapis.com/elasticsearch-cuvs-snapshots");
+            });
+            repos.exclusiveContent(exclusiveRepo -> {
+                exclusiveRepo.filter(descriptor -> descriptor.includeVersionByRegex("com\\.nvidia\\.cuvs", ".*", ".*-SNAPSHOT"));
+                exclusiveRepo.forRepositories(cuvsRepo);
+            });
+        }
     }
 }

build-tools-internal/version.properties

@@ -19,7 +19,8 @@ commons_lang3     = 3.9
 google_oauth_client = 1.34.1
 awsv2sdk            = 2.31.78
 reactive_streams    = 1.0.4
-ldapsdk           = 7.0.3
+cuvs_java           = 25.10.0
+ldapsdk             = 7.0.3
 
 antlr4            = 4.13.1
 # bouncy castle version for non-fips. fips jars use a different version

libs/gpu-codec/build.gradle

@@ -16,7 +16,7 @@ dependencies {
   api project(':libs:logging')
   compileOnly "org.apache.lucene:lucene-core:${versions.lucene}"
   compileOnly project(':server')
-  implementation('com.nvidia.cuvs:cuvs-java:25.10.0')
+  implementation("com.nvidia.cuvs:cuvs-java:${versions.cuvs_java}")
 
   testImplementation(project(":test:framework")) {
     exclude group: 'org.elasticsearch', module: 'gpu-codec'

x-pack/plugin/gpu/build.gradle

@@ -20,7 +20,7 @@ dependencies {
   testImplementation(testArtifact(project(xpackModule('core'))))
   testImplementation(testArtifact(project(':server')))
   yamlRestTestImplementation project(':libs:gpu-codec')
-  internalClusterTestImplementation('com.nvidia.cuvs:cuvs-java:25.10.0')
+  internalClusterTestImplementation("com.nvidia.cuvs:cuvs-java:${versions.cuvs_java}")
 
   clusterModules project(xpackModule('gpu'))
 }