From ca83b129b993d80d743cb797e4524f34fa135576 Mon Sep 17 00:00:00 2001
From: Quentin Bourgerie <bourgerie.quentin@gmail.com>
Date: Fri, 6 Jan 2023 17:53:23 +0100
Subject: [PATCH] chore(ci): fix gpu build on aws using docker

---
 .github/workflows/aws_build.yml | 13 +++++--------
 ci/slab.toml                    |  2 +-
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/aws_build.yml b/.github/workflows/aws_build.yml
index 2fd2eac8f..2aa476da2 100644
--- a/.github/workflows/aws_build.yml
+++ b/.github/workflows/aws_build.yml
@@ -28,6 +28,7 @@ env:
   RESULTS_FILENAME: parsed_benchmark_results_${{ github.sha }}.json
   MAKE_RULE_TESTS: run-tests run-end-to-end-dataflow-tests
   DOCKER_IMAGE_TEST: ghcr.io/zama-ai/concrete-compiler
+  DOCKER_GPU_OPTION:
   CUDA_SUPPORT: OFF
   CUDA_PATH: /usr/local/cuda-11.7
   GCC_VERSION: 11
@@ -67,6 +68,7 @@ jobs:
         run: |
           echo "CUDA_SUPPORT=ON" >> "${GITHUB_ENV}"
           echo "MAKE_RULE_TESTS=run-end-to-end-tests-gpu" >> "${GITHUB_ENV}"
+          echo "DOCKER_GPU_OPTION=--gpus all" >> "${GITHUB_ENV}"
 
       # Free 4Gb of workspace
       - name: Freeing space
@@ -97,7 +99,6 @@ jobs:
       - name: Create build dir
         run: mkdir build
 
-      # FIXME: GPU is not supported yet, docker image has to be rebuilt with CUDA first.
       - name: Build compiler
         uses: addnab/docker-run-action@v3
         id: build-compiler
@@ -112,15 +113,9 @@ jobs:
             -v ${{ github.workspace }}/build:/build
             -v ${{ env.SSH_AUTH_SOCK }}:/ssh.socket
             -e SSH_AUTH_SOCK=/ssh.socket
+            ${{ env.DOCKER_GPU_OPTION }}
           shell: bash
           run: |
-            export CUDA_PATH=$CUDA_PATH            
-            export LD_LIBRARY_PATH=$CUDA_PATH/lib:$LD_LIBRARY_PATH
-            export CC=$(which gcc)
-            export CXX=$(which g++)
-            export CUDAHOSTCXX=$CXX
-            export CUDACXX=$CUDA_PATH/bin/nvcc
-          
             set -e
             cd /compiler
             rm -rf /build/*
@@ -153,11 +148,13 @@ jobs:
             -v ${{ github.workspace }}/compiler:/compiler
             -v ${{ github.workspace }}/KeySetCache:/tmp/KeySetCache
             -v ${{ github.workspace }}/build:/build
+            ${{ env.DOCKER_GPU_OPTION }}
           shell: bash
           run: |
             set -e
             cd /compiler
             pip install pytest
+            mkdir -p /tmp/concrete_compiler/gpu_tests/
             make DATAFLOW_EXECUTION_ENABLED=ON CCACHE=ON Python3_EXECUTABLE=$PYTHON_EXEC BUILD_DIR=/build ${{ env.MAKE_RULE_TESTS }}
             chmod -R ugo+rwx /tmp/KeySetCache
 
diff --git a/ci/slab.toml b/ci/slab.toml
index e7fa46146..2b2fcbf00 100644
--- a/ci/slab.toml
+++ b/ci/slab.toml
@@ -14,7 +14,7 @@ security_group= ["sg-0bf1c1d79c97bc88f", ]
 
 [profile.gpu]
 region = "us-east-1"
-image_id = "ami-03deb184ab492226b"
+image_id = "ami-0c4773f5626d919b6"
 instance_type = "p3.2xlarge"
 subnet_id = "subnet-8123c9e7"
 security_group= ["sg-0f8b52622a2669491", ]