From 6726475db2c907262fe7ab423d90fde1998cd39f Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@localai.io>
Date: Wed, 18 Oct 2023 21:19:15 +0200
Subject: [PATCH] split extra/core images

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 .github/workflows/test-gpu.yml |  2 +-
 Makefile                       |  4 +++-
 tests/e2e/e2e_test.go          | 15 ++++++++++-----
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/test-gpu.yml b/.github/workflows/test-gpu.yml
index ee08b9332e4a..9a1cf8bb14c4 100644
--- a/.github/workflows/test-gpu.yml
+++ b/.github/workflows/test-gpu.yml
@@ -71,7 +71,7 @@ jobs:
           if [ ! -e /run/systemd/system ]; then
             sudo mkdir /run/systemd/system
           fi
-          make prepare-e2e run-e2e-image test-e2e
+          make BUILD_TYPE=cublas prepare-e2e run-e2e-image test-e2e
       - name: Release space from worker ♻
         if: always()
         run: |
diff --git a/Makefile b/Makefile
index 716b8edaab0d..6986c892b7d0 100644
--- a/Makefile
+++ b/Makefile
@@ -344,7 +344,9 @@ run-e2e-image:
 
 test-e2e:
 	@echo 'Running e2e tests'
-	LOCALAI_API=http://$(E2E_BRIDGE_IP):5390/v1 $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts 5 -v -r ./tests/e2e
+	BUILD_TYPE=$(BUILD_TYPE) \
+	LOCALAI_API=http://$(E2E_BRIDGE_IP):5390/v1 \
+	$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts 5 -v -r ./tests/e2e
 
 teardown-e2e:
 	rm -rf ./tests/e2e-fixtures/ggllm-test-model.bin
diff --git a/tests/e2e/e2e_test.go b/tests/e2e/e2e_test.go
index f39b6cd599bf..dbf6d66ba323 100644
--- a/tests/e2e/e2e_test.go
+++ b/tests/e2e/e2e_test.go
@@ -3,6 +3,7 @@ package e2e_test
 import (
 	"context"
 	"fmt"
+	"os"
 	"os/exec"
 
 	. "github.com/onsi/ginkgo/v2"
@@ -35,11 +36,15 @@ var _ = Describe("E2E test", func() {
 		// Check that the GPU was used
 		AfterEach(func() {
 			// Execute docker logs $$(docker ps -q --filter ancestor=localai-tests) as a command and check the output
-			cmd := exec.Command("/bin/bash", "-xce", "docker logs $(docker ps -q --filter ancestor=localai-tests)")
-			out, err := cmd.CombinedOutput()
-			Expect(err).ToNot(HaveOccurred(), string(out))
-			Expect(string(out)).To(ContainSubstring("found 1 CUDA devices"), string(out))
-			Expect(string(out)).To(ContainSubstring("using CUDA for GPU acceleration"), string(out))
+			if os.Getenv("BUILD_TYPE") == "cublas" {
+				cmd := exec.Command("/bin/bash", "-xce", "docker logs $(docker ps -q --filter ancestor=localai-tests)")
+				out, err := cmd.CombinedOutput()
+				Expect(err).ToNot(HaveOccurred(), string(out))
+				Expect(string(out)).To(ContainSubstring("found 1 CUDA devices"), string(out))
+				Expect(string(out)).To(ContainSubstring("using CUDA for GPU acceleration"), string(out))
+			} else {
+				fmt.Println("Skipping GPU check")
+			}
 		})
 
 		Context("Generates text", func() {