Build GPU docker image on linux if nvidia-smi exists (#61)

beam-cloud · Jan 22, 2024 · 22dc23e · 22dc23e
1 parent b653d89
commit 22dc23e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/Makefile b/Makefile
@@ -27,7 +27,7 @@ beam-worker:
 	bin/delete_workers.sh
 
 beam-runner:
-	for target in py312 py311 py310 py39 py38; do \
+	for target in py311 py310 py39 py38; do \
 		docker build . --target $$target --platform=linux/amd64 -f ./docker/Dockerfile.runner -t localhost:5001/beam-runner:$$target-$(runnerTag); \
 		docker push localhost:5001/beam-runner:$$target-$(runnerTag); \
 	done

diff --git a/bin/k3d.sh b/bin/k3d.sh
@@ -16,8 +16,12 @@ k3d_up() {
   os_type="$(uname)"
   case "$os_type" in
     Linux*)
-      check_gpu_linux && extra_args="--gpus=all --image=localhost:5001/rancher/k3s:latest" || extra_args=""
-      docker build . -f ./docker/Dockerfile.k3d -t localhost:5001/rancher/k3s:latest
+      if check_gpu_linux; then
+        extra_args="--gpus=all --image=localhost:5001/rancher/k3s:latest"
+        docker build . -f ./docker/Dockerfile.k3d -t localhost:5001/rancher/k3s:latest
+      else
+        extra_args=""
+      fi
       ;;
     Darwin*)
       extra_args=""