From 66ec5cda55e063bb81e570a9af7868edf294dd7e Mon Sep 17 00:00:00 2001
From: Benedikt Daurer <benedikt.daurer@diamond.ac.uk>
Date: Wed, 31 Jul 2024 15:53:01 +0100
Subject: [PATCH] add argument for CUDA version and reformat

---
 Dockerfile | 33 ++++++++++++++++++++++++++++-----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index c16f2062c..8857df070 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -4,6 +4,9 @@ ARG MPI=openmpi
 # Select Platform: core, full, pycuda or cupy
 ARG PLATFORM=cupy
 
+# Select CUDA version
+ARG CUDAVERSION=12.4
+
 # Pull from mambaforge and install XML and ssh
 FROM condaforge/mambaforge as base
 ENV DEBIAN_FRONTEND=noninteractive
@@ -26,18 +29,24 @@ RUN mamba env update -n base -f dependencies.yml
 
 # Pull from MPI build and install accelerate/pycuda dependencies
 FROM mpi as pycuda
+ARG CUDAVERSION
 COPY ./ptypy/accelerate/cuda_pycuda/dependencies.yml ./dependencies.yml
 COPY ./cufft/dependencies.yml ./dependencies_cufft.yml
-RUN mamba env update -n base -f dependencies.yml && mamba env update -n base -f dependencies_cufft.yml
+RUN mamba install cuda-version=${CUDAVERSION} && \
+    mamba env update -n base -f dependencies.yml && \
+    mamba env update -n base -f dependencies_cufft.yml
 
 # Pull from MPI build and install accelerate/cupy dependencies
 FROM mpi as cupy
+ARG CUDAVERSION
 COPY ./ptypy/accelerate/cuda_cupy/dependencies.yml ./dependencies.yml
 COPY ./cufft/dependencies.yml ./dependencies_cufft.yml
-RUN mamba env update -n base -f dependencies.yml && mamba env update -n base -f dependencies_cufft.yml
+RUN mamba install cuda-version=${CUDAVERSION} && \
+    mamba env update -n base -f dependencies.yml && \
+    mamba env update -n base -f dependencies_cufft.yml
 
-# Pull from existing image with ptypy dependencies and set up testing
-FROM ${PLATFORM} as runtime
+# Pull from platform specific image and install ptypy 
+FROM ${PLATFORM} as build
 COPY pyproject.toml setup.py ./
 COPY ./scripts ./scripts
 COPY ./templates ./templates
@@ -45,7 +54,21 @@ COPY ./benchmark ./benchmark
 COPY ./cufft ./cufft
 COPY ./ptypy ./ptypy
 RUN pip install .
-RUN if [ "$PLATFORM" = "pycuda" ] || [ "$PLATFORM" = "cupy" ] ; then pip install ./cufft ; fi
+
+# For core/full build, no post processing needed
+FROM build as core-post
+FROM build as full-post
+
+# For pycuda build, install filtered cufft
+FROM build as pycuda-post
+RUN pip install ./cufft
+
+# For pycuda build, install filtered cufft
+FROM build as cupy-post
+RUN pip install ./cufft
+
+# Platform specific runtime container
+FROM ${PLATFORM}-post as runtime
 
 # Run PtyPy run script as entrypoint
 ENTRYPOINT ["ptypy.run"]