From 24934358f5c539d65ef0d80377640d4b155d7138 Mon Sep 17 00:00:00 2001
From: Ana Espinoza <respinoza@ucar.edu>
Date: Tue, 30 Jan 2024 18:24:40 -0700
Subject: [PATCH 1/2] Docker image improvements

---
 gpu/nvidia-tensorflow-jupyterhub/Dockerfile   | 59 +++++++++----------
 .../environment-tf.yml                        | 14 +++++
 .../fix-permissions                           | 35 -----------
 3 files changed, 41 insertions(+), 67 deletions(-)
 create mode 100644 gpu/nvidia-tensorflow-jupyterhub/environment-tf.yml
 delete mode 100644 gpu/nvidia-tensorflow-jupyterhub/fix-permissions

diff --git a/gpu/nvidia-tensorflow-jupyterhub/Dockerfile b/gpu/nvidia-tensorflow-jupyterhub/Dockerfile
index 2732d45..ddd9bf3 100644
--- a/gpu/nvidia-tensorflow-jupyterhub/Dockerfile
+++ b/gpu/nvidia-tensorflow-jupyterhub/Dockerfile
@@ -1,36 +1,31 @@
-FROM nvcr.io/nvidia/tensorflow:22.04-tf2-py3
-RUN pip install --no-cache-dir jupyterhub==3.0.0
-RUN pip install --no-cache-dir --upgrade jupyterlab notebook jupyter_server
-
-# Initialization of user copied from:
-# https://github.com/jupyter/docker-stacks/blob/main/docker-stacks-foundation/Dockerfile
-
-ARG NB_USER="jovyan"
-ARG NB_UID="1000"
-ARG NB_GID="100"
-
-# Fix: https://github.com/hadolint/hadolint/wiki/DL4006
-# Fix: https://github.com/koalaman/shellcheck/wiki/SC3014
-SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+ARG BASE_CONTAINER=jupyter/minimal-notebook
+FROM $BASE_CONTAINER
 
+# Install any additional system packages
 USER root
-
-COPY fix-permissions /usr/local/bin/fix-permissions
-RUN chmod a+rx /usr/local/bin/fix-permissions
-
-ENV HOME="/home/${NB_USER}"
-
-# Enable prompt color in the skeleton .bashrc before creating the default NB_USER
-# hadolint ignore=SC2016
-RUN sed -i 's/^#force_color_prompt=yes/force_color_prompt=yes/' /etc/skel/.bashrc
-
-# Create NB_USER with name jovyan user with UID=1000 and in the 'users' group
-# and make sure these dirs are writable by the `users` group.
-RUN echo "auth requisite pam_deny.so" >> /etc/pam.d/su && \
-    useradd -l -m -s /bin/bash -N -u "${NB_UID}" "${NB_USER}" && \
-    chmod g+w /etc/passwd && \
-    fix-permissions "${HOME}"
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends vim curl && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+USER $NB_UID
+
+# Install some packages into base, create GPU enabled tensorflow conda
+# environment, and clean
+ARG ENVIRONMENT_FILE=environment-tf.yml
+ADD $ENVIRONMENT_FILE /tmp/environment.yml
+RUN mamba install --quiet --yes \
+      'conda-forge::nb_conda_kernels' \
+      'conda-forge::jupyterlab-git' \
+      'conda-forge::ipywidgets' && \
+    mamba env update -f /tmp/environment.yml && \
+    pip install --no-cache-dir nbgitpuller && \
+    mamba clean --all -f -y && \
+    jupyter lab clean -y && \
+    npm cache clean --force && \
+    rm -rf /home/$NB_USER/.cache/yarn && \
+    rm -rf /home/$NB_USER/.node-gyp && \
+    fix-permissions $CONDA_DIR && \
+    fix-permissions /home/$NB_USER
 
 USER ${NB_UID}
-
-WORKDIR "${HOME}"
diff --git a/gpu/nvidia-tensorflow-jupyterhub/environment-tf.yml b/gpu/nvidia-tensorflow-jupyterhub/environment-tf.yml
new file mode 100644
index 0000000..24fb696
--- /dev/null
+++ b/gpu/nvidia-tensorflow-jupyterhub/environment-tf.yml
@@ -0,0 +1,14 @@
+name: tensorflow
+channels:
+ - conda-forge
+dependencies:
+  # Required by JupyterLab
+  - python=3
+  - nb_conda_kernels
+  - ipykernel
+  - pip:
+    # Tensorflow dependencies
+    # The [and-cuda] extra is necessary to enable GPU capability
+    # If something doesn't seem right, ensure that your tensorflow, cuda, and
+    # python version match: https://www.tensorflow.org/install/source#gpu
+    - tensorflow[and-cuda]
diff --git a/gpu/nvidia-tensorflow-jupyterhub/fix-permissions b/gpu/nvidia-tensorflow-jupyterhub/fix-permissions
deleted file mode 100644
index d167578..0000000
--- a/gpu/nvidia-tensorflow-jupyterhub/fix-permissions
+++ /dev/null
@@ -1,35 +0,0 @@
-#!/bin/bash
-# set permissions on a directory
-# after any installation, if a directory needs to be (human) user-writable,
-# run this script on it.
-# It will make everything in the directory owned by the group ${NB_GID}
-# and writable by that group.
-# Deployments that want to set a specific user id can preserve permissions
-# by adding the `--group-add users` line to `docker run`.
-
-# uses find to avoid touching files that already have the right permissions,
-# which would cause massive image explosion
-
-# right permissions are:
-# group=${NB_GID}
-# AND permissions include group rwX (directory-execute)
-# AND directories have setuid,setgid bits set
-
-set -e
-
-for d in "$@"; do
-    find "${d}" \
-        ! \( \
-            -group "${NB_GID}" \
-            -a -perm -g+rwX \
-        \) \
-        -exec chgrp "${NB_GID}" -- {} \+ \
-        -exec chmod g+rwX -- {} \+
-    # setuid, setgid *on directories only*
-    find "${d}" \
-        \( \
-            -type d \
-            -a ! -perm -6000 \
-        \) \
-        -exec chmod +6000 -- {} \+
-done

From 4f1090495d973c7a2042025f51f944489c91456a Mon Sep 17 00:00:00 2001
From: Ana Espinoza <respinoza@ucar.edu>
Date: Tue, 30 Jan 2024 18:26:52 -0700
Subject: [PATCH 2/2] Initial try at a minimal pytorch environment

---
 gpu/nvidia-tensorflow-jupyterhub/Makefile     | 19 +++++++++++++++++++
 .../environment-pytorch.yml                   | 19 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 gpu/nvidia-tensorflow-jupyterhub/environment-pytorch.yml

diff --git a/gpu/nvidia-tensorflow-jupyterhub/Makefile b/gpu/nvidia-tensorflow-jupyterhub/Makefile
index cd3bc1e..0f9a7fd 100644
--- a/gpu/nvidia-tensorflow-jupyterhub/Makefile
+++ b/gpu/nvidia-tensorflow-jupyterhub/Makefile
@@ -17,3 +17,22 @@ shell:
 
 tag:
 	git tag -a ${VERSION} -m "Tag ${VERSION} in Github repository"
+
+# ####################################################
+# For optionally building with Pytorch instead of TF
+# ####################################################
+
+PYTORCH = zonca/pytorch-jupyterhub
+
+pytorch:
+	docker build -t $(PYTORCH):${PT_VERSION} --build-arg="ENVIRONMENT_FILE=environment-pytorch.yml" .
+
+cleanpytorch:
+	docker build --no-cache -t $(PYTORCH):${PT_VERSION} --build-arg="ENVIRONMENT_FILE=environment-pytorch.yml" .
+
+pushpytorch:
+	docker push $(PYTORCH):${PT_VERSION}
+
+shellpytorch:
+	docker run -it $(PYTORCH):${PT_VERSION} bash
+
diff --git a/gpu/nvidia-tensorflow-jupyterhub/environment-pytorch.yml b/gpu/nvidia-tensorflow-jupyterhub/environment-pytorch.yml
new file mode 100644
index 0000000..45dab0b
--- /dev/null
+++ b/gpu/nvidia-tensorflow-jupyterhub/environment-pytorch.yml
@@ -0,0 +1,19 @@
+name: pytorch
+channels:
+  - conda-forge
+  # Begin pytorch channels  
+  - pytorch
+  - nvidia
+  # End pytorch channels  
+dependencies:
+  # Required by JupyterLab
+  - python=3
+  - nb_conda_kernels
+  - ipykernel
+  # pytorch
+  # Environment file adapted from the instruction found at:
+  # https://pytorch.org/get-started/locally/
+  - pytorch
+  - torchvision
+  - torchaudio
+  - pytorch-cuda=12.1