From 53cec287a5e23009d5d0e94937254b9103761eeb Mon Sep 17 00:00:00 2001
From: "Joshua L. Phillips" <Joshua.Phillips@mtsu.edu>
Date: Tue, 20 Feb 2024 11:59:10 -0600
Subject: [PATCH] Bump to 2024.2 and add PyTorch 2.2.0

---
 Dockerfile | 49 +++++++++++++++++++++++++++++++++++--------------
 1 file changed, 35 insertions(+), 14 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 7a8193f..acc7f28 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,7 +1,7 @@
-FROM quay.io/jupyter/minimal-notebook:2023-11-29 AS build-base
+FROM quay.io/jupyter/minimal-notebook:2024-02-13 AS build-base
 
 LABEL maintainer="Joshua L. Phillips <https://www.cs.mtsu.edu/~jphillips/>"
-LABEL release-date="2023-12-04"
+LABEL release-date="2024-02-20"
 
 USER root
 
@@ -50,9 +50,14 @@ RUN mamba install --yes \
     mamba clean --all -f -y
 
 FROM build-base AS build-qiime2
-ARG Q2_VERSION=2023.9
+ARG Q2_VERSION=2024.2
 ARG Q2_DISTRIBUTION=amplicon
 ENV QIIME2_ENV=qiime2-$Q2_DISTRIBUTION-$Q2_VERSION
+ENV XLA_FLAGS="--xla_gpu_cuda_data_dir=/opt/conda/envs/qiime2-$Q2_DISTRIBUTION-$Q2_VERSION/pkgs/cuda-toolkit"
+
+USER root
+RUN echo -e "/opt/conda/envs/${QIIME2_ENV}/lib" | tee /etc/ld.so.conf.d/zzz-conda.conf
+USER ${NB_UID}
 
 # Qiime2
 RUN wget https://data.qiime2.org/distro/$Q2_DISTRIBUTION/$QIIME2_ENV-py38-linux-conda.yml && \
@@ -64,8 +69,16 @@ RUN mamba run -n ${QIIME2_ENV} python -m ipykernel install --name ${QIIME2_ENV}
 # SHELL ["mamba", "run", "-n", "${QIIME2_ENV}", "/bin/bash", "-o", "pipefail", "-c"]
 # SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 
-FROM build-qiime2 AS build-tensorflow
-# GPU-enabled stack
+# RESCRIPt - we use this for much of our work
+# even if it (was) not in the Qiime2 amplicon
+# distribution by default, so added here (no longer needed).
+# RUN mamba install -n ${QIIME2_ENV} -c conda-forge -c bioconda -c qiime2 \
+#     -c https://packages.qiime2.org/qiime2/$Q2_VERSION/shotgun/released/ -c defaults \
+#     xmltodict "q2-types-genomics>$Q2_VERSION" ncbi-datasets-pylib && \
+#     mamba run -n ${QIIME2_ENV} pip install --no-cache-dir git+https://github.com/bokulich-lab/RESCRIPt.git
+
+FROM build-qiime2 AS build-deeplearning
+# GPU-enabled Tensorflow
 RUN CONDA_OVERRIDE_CUDA="11.8" mamba install -n $QIIME2_ENV --yes \
     bokeh \
     cudatoolkit==11.8.0 \
@@ -79,6 +92,22 @@ RUN CONDA_OVERRIDE_CUDA="11.8" mamba install -n $QIIME2_ENV --yes \
     && \
     mamba clean --all -f -y
 
+# PyTorch
+RUN mamba run -n ${QIIME2_ENV} pip install --no-cache-dir \
+    deepspeed \
+    jsonargparse[signatures] \
+    lightning \
+    tokenizers \
+    torch \
+    torchaudio \
+    torchinfo \
+    torchmetrics \
+    torchtext \
+    torchview \
+    torchvision \
+    transformers \
+    --extra-index-url https://download.pytorch.org/whl/cu118
+
 # Note that we cannot clean mamba -at all- after
 # this point - or it will remove the dev tools from the stack...
 RUN mamba install -n $QIIME2_ENV --yes \
@@ -89,15 +118,7 @@ RUN mamba install -n $QIIME2_ENV --yes \
 # Configure user environment (needed for singularity/apptainer)
 RUN cp /etc/skel/.bash_logout /etc/skel/.bashrc /etc/skel/.profile /home/${NB_USER}/. && conda init && mamba init
 
-FROM build-tensorflow AS build-deep-dna
-
-# RESCRIPt - we use this for much of our work
-# even if it's not in the Qiime2 amplicon
-# distribution by default, so added here.
-RUN mamba install -n ${QIIME2_ENV} -c conda-forge -c bioconda -c qiime2 \
-    -c https://packages.qiime2.org/qiime2/2023.9/shotgun/released/ -c defaults \
-    xmltodict 'q2-types-genomics>2023.2' ncbi-datasets-pylib && \
-    mamba run -n ${QIIME2_ENV} pip install --no-cache-dir git+https://github.com/bokulich-lab/RESCRIPt.git
+FROM build-deeplearning AS build-deep-dna
 
 # SetBERT Packages
 RUN git clone https://github.com/DLii-Research/deep-dna && \