Add job for testing TF integrations on GPU (#9979)

Part of #9855
iree-org · Aug 2, 2022 · c394dbb · c394dbb
1 parent 76bf2f3
commit c394dbb
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 6 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -258,8 +258,10 @@ jobs:
             --gpus all \
             --env NVIDIA_DRIVER_CAPABILITIES=all \
             gcr.io/iree-oss/nvidia@sha256:7c2f56db65e656c15e6c96b5812a8275dd53c82bf41221192f9ba8a451aad870 \
-            bash -c \
-            "./build_tools/scripts/check_cuda.sh && ./build_tools/scripts/check_vulkan.sh && ./build_tools/cmake/ctest_all.sh ${BUILD_DIR}"
+            bash -euo pipefail -c \
+              "./build_tools/scripts/check_cuda.sh
+              ./build_tools/scripts/check_vulkan.sh
+              ./build_tools/cmake/ctest_all.sh ${BUILD_DIR}"
 
   test_tf_integrations:
     needs: [build_tf_integrations, build_all]
@@ -299,6 +301,50 @@ jobs:
             build_tools/cmake/run_tf_tests.sh \
             "${BUILD_DIR}"
 
+  test_tf_integrations_gpu:
+    needs: [build_tf_integrations, build_all]
+    if: github.event_name != 'pull_request'
+    runs-on:
+      # Hacks, and order matters. See the comment at the top of the file.
+      - self-hosted
+      - runner-group=${{ github.event_name == 'pull_request' && 'presubmit' || 'postsubmit' }}
+      - gpu
+      - os-family=Linux
+    env:
+      BUILD_DIR: ${{ needs.build_all.outputs.build-dir }}
+      BUILD_DIR_ARCHIVE: ${{ needs.build_all.outputs.build-dir-archive }}
+      BUILD_DIR_GCS_ARTIFACT: ${{ needs.build_all.outputs.gcs-artifact }}
+      TF_BINARIES_DIR: ${{ needs.build_tf_integrations.outputs.binaries-dir }}
+      TF_BINARIES_ARCHIVE: ${{ needs.build_tf_integrations.outputs.binaries-archive }}
+      TF_BINARIES_GCS_ARTIFACT: ${{ needs.build_tf_integrations.outputs.gcs-artifact }}
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@7884fcad6b5d53d10323aee724dc68d8b9096a2e # v2
+        with:
+          submodules: true
+      - name: "Downloading TF binaries archive"
+        run: gcloud alpha storage cp "${TF_BINARIES_GCS_ARTIFACT}" "${TF_BINARIES_ARCHIVE}"
+      - name: "Extracting TF binaries archive"
+        run: tar -xf "${TF_BINARIES_ARCHIVE}"
+      - name: "Symlinking TF binaries"
+        run: |
+          ./integrations/tensorflow/symlink_binaries.sh "$(realpath "${TF_BINARIES_DIR}")"
+      - name: "Downloading build dir archive"
+        run: gcloud alpha storage cp "${BUILD_DIR_GCS_ARTIFACT}" "${BUILD_DIR_ARCHIVE}"
+      - name: "Extracting build dir archive"
+        run: tar -xf "${BUILD_DIR_ARCHIVE}"
+      - name: "Running TF integrations tests"
+        run: |
+          ./build_tools/github_actions/docker_run.sh \
+            --env IREE_LLVM_CPU_DISABLE=1 \
+            --gpus all \
+            --env NVIDIA_DRIVER_CAPABILITIES=all \
+            gcr.io/iree-oss/frontends-nvidia@sha256:e934ed09e9e60c28ebe11a02f37a993dd975db40118d410c4279d0fa2d4e6b9a \
+            bash -euo pipefail -c \
+              "./build_tools/scripts/check_cuda.sh
+              ./build_tools/scripts/check_vulkan.sh
+              build_tools/cmake/run_tf_tests.sh ${BUILD_DIR}"
+
   asan:
     runs-on:
       # Hacks, and order matters. See the comment at the top of the file.
@@ -367,7 +413,7 @@ jobs:
             --env "BUILD_RISCV_DIR=${BUILD_RISCV_DIR}" \
             --env "IREE_HOST_BINARY_ROOT=${HOST_BINARY_ROOT}" \
             gcr.io/iree-oss/riscv@sha256:720bc0215d8462ea14352edc22710a6ce4c0c1daff581d179dd173885f1d8a35 \
-            bash -c \
+            bash -euo pipefail -c \
             "./build_tools/cmake/build_riscv.sh && tests/riscv32/smoke.sh"
 
   riscv64:
@@ -421,7 +467,7 @@ jobs:
             --env "IREE_IMPORT_TFLITE_BIN=${TF_BINARIES_DIR}/iree-import-tflite" \
             --env "LLVM_BIN_DIR=${BUILD_DIR}/third_party/llvm-project/llvm/bin" \
             gcr.io/iree-oss/riscv@sha256:720bc0215d8462ea14352edc22710a6ce4c0c1daff581d179dd173885f1d8a35 \
-            bash -c \
+            bash -euo pipefail -c \
               "./build_tools/cmake/build_riscv.sh && ./tests/riscv64/smoke.sh"
 
 
@@ -434,6 +480,8 @@ jobs:
       - test_all
       - test_gpu
       - host_tools_assertions
+      - test_tf_integrations
+      - test_tf_integrations_gpu
       - asan
       - tsan
       - riscv32

diff --git a/build_tools/cmake/run_tf_tests.sh b/build_tools/cmake/run_tf_tests.sh
@@ -13,6 +13,7 @@ cd "${ROOT_DIR}"
 
 BUILD_DIR="$1"
 IREE_VULKAN_DISABLE="${IREE_VULKAN_DISABLE:-0}"
+IREE_LLVM_CPU_DISABLE="${IREE_LLVM_CPU_DISABLE:-0}"
 
 source "${BUILD_DIR}/.env" && export PYTHONPATH
 
@@ -32,6 +33,10 @@ if (( ${IREE_VULKAN_DISABLE} != 1 )); then
   CMD+=(-D FEATURES=vulkan)
 fi
 
+if (( ${IREE_LLVM_CPU_DISABLE} == 1 )); then
+  CMD+=(-D DISABLE_FEATURES=llvmcpu)
+fi
+
 if "${CMD[@]}"; then
   tests_passed=1
 else

diff --git a/build_tools/scripts/check_cuda.sh b/build_tools/scripts/check_cuda.sh
@@ -8,7 +8,7 @@
 
 # Checks that CUDA is working correctly and logs some useful information
 
-set -x
+set -xeuo pipefail
 
 dpkg -l | grep nvidia
 nvidia-smi
diff --git a/build_tools/scripts/check_vulkan.sh b/build_tools/scripts/check_vulkan.sh
@@ -8,7 +8,7 @@
 
 # Checks that Vulkan is working correctly and logs some useful information
 
-set -x
+set -xeuo pipefail
 
 # Print Vulkan related information: SDK version and GPU ICD version
 vulkaninfo 2> /tmp/vulkaninfo.stderr 1> /tmp/vulkaninfo.stdout