Merge branch 'main' into pixi-lockfiles

Author: marcelovilla

.github/workflows/build-push-docker.yaml

@@ -91,6 +91,7 @@ jobs:
           echo "GPU Platform Matrix"
           echo "BASE_IMAGE=$GPU_BASE_IMAGE" >> $GITHUB_ENV
           echo "IMAGE_SUFFIX=-$GPU_IMAGE_SUFFIX" >> $GITHUB_ENV
+          echo "GPU=true" >> $GITHUB_ENV
 
       - name: "Generate Docker images tags 🏷️"
         id: meta
@@ -123,5 +124,7 @@ jobs:
           labels: ${{ steps.meta.outputs.labels }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
-          build-args: BASE_IMAGE=${{ env.BASE_IMAGE }}
+          build-args: |
+            BASE_IMAGE=${{ env.BASE_IMAGE }}
+            GPU=${{ env.GPU }}
           platforms: linux/amd64,linux/arm64

Dockerfile

@@ -1,4 +1,5 @@
-FROM ubuntu:24.04 AS builder
+ARG BASE_IMAGE=ubuntu:24.04
+FROM ${BASE_IMAGE} AS builder
 LABEL MAINTAINER="Nebari development team"
 
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
@@ -22,6 +23,11 @@ ENV DEFAULT_ENV=default
 # ========== dask-worker install ===========
 FROM builder AS dask-worker
 
+ARG GPU
+ENV LD_LIBRARY_PATH=${GPU:+/usr/local/nvidia/lib64}
+ENV NVIDIA_PATH=${GPU:+/usr/local/nvidia/bin}
+ENV PATH=${GPU:+/usr/local/nvidia/bin:}${PATH}
+
 COPY --from=builder /root/.pixi /root/.pixi
 COPY dask-worker/pixi.toml /opt/dask-worker/pixi.toml
 COPY dask-worker/pixi.lock /opt/dask-worker/pixi.lock
@@ -86,6 +92,15 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
 # ========== jupyterlab install ===========
 FROM intermediate AS jupyterlab
 
+ARG GPU
+
+ENV CONDA_DIR=/opt/conda \
+    DEFAULT_ENV=default \
+    LD_LIBRARY_PATH=${GPU:+/usr/local/nvidia/lib64} \
+    NVIDIA_PATH=${GPU:+/usr/local/nvidia/bin}
+
+ENV PATH=${GPU:+/usr/local/nvidia/bin:}${PATH}
+
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get update && apt-get install -y --no-install-recommends \