From 567fc9cd99bf3ffdb24b05e9471b1ce88e9573f7 Mon Sep 17 00:00:00 2001
From: Daphne Cornelisse <cor.daphne@gmail.com>
Date: Tue, 21 May 2024 20:17:45 -0400
Subject: [PATCH] Add Dockerfile and launsh.sh for HPC

---
 hpc/Dockerfile | 18 +++++++++++++++
 hpc/launch.sh  | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 hpc/Dockerfile
 create mode 100644 hpc/launch.sh

diff --git a/hpc/Dockerfile b/hpc/Dockerfile
new file mode 100644
index 00000000..bafdf861
--- /dev/null
+++ b/hpc/Dockerfile
@@ -0,0 +1,18 @@
+FROM continuumio/miniconda3:24.1.2-0
+
+# Set environment variables
+ENV WANDB_CACHE_DIR=$WANDB_CACHE_DIR
+ENV CONDA_ENVS_DIRS=$CONDA_ENVS_DIRS
+
+# Install dependencies including cmake >=3.20
+RUN apt-get update && apt-get install -y \
+    software-properties-common \
+    apt-transport-https \
+    gpg-agent \
+    wget \
+    && wget -qO - https://apt.kitware.com/keys/kitware-archive-latest.asc | gpg --dearmor -o /usr/share/keyrings/kitware-archive-keyring.gpg \
+    && echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ bionic main' > /etc/apt/sources.list.d/kitware.list \
+    && apt-get update \
+    && apt-get install -y cmake git \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
diff --git a/hpc/launch.sh b/hpc/launch.sh
new file mode 100644
index 00000000..c3599ea0
--- /dev/null
+++ b/hpc/launch.sh
@@ -0,0 +1,61 @@
+# Script for launching the singularity image for gpudrive
+# ---
+# Note: Script must be run from the gpudrive repository (i.e. /scratch/$USER/gpudrive/)
+# and executed each time you want to use gpudrive (lab version) on the HPC (e.g. Greene).
+
+# Constants
+PROJECT="gpudrive"
+PROJECT_DOCKER=docker://daphnecor/gpudrive
+SINGULARITY_IMAGE=./hpc/gpudrive.sif
+OVERLAY_LOC=/scratch/work/public/overlay-fs-ext3
+OVERLAY_FILE=overlay-15GB-500K.ext3
+
+# Overwrite wandb cache dir to avoid storage capacity problems
+WANDB_CACHE_DIR='${WANDB_CACHE_DIR}'
+
+# Check if singularity image exists, if not pull Singularity image from Docker Hub
+if [ ! -f "${SINGULARITY_IMAGE}" ]; then
+    echo "Pulling Docker container from ${PROJECT_DOCKER}"
+    singularity pull $SINGULARITY_IMAGE $PROJECT_DOCKER
+fi
+
+# Check if overlay file exists, if not create it
+if [ ! -f "./hpc/${OVERLAY_FILE}" ]; then  # Overlay file does not exist
+    echo "Setting up ${PROJECT_DOCKER} with initial overlay ${OVERLAY_FILE}.gz"
+
+    if [ ! -f "${OVERLAY_FILE}.gz" ]; then  # Overlay file has not been copiepd yet
+        echo "Copying overlay ${OVERLAY_FILE}.gz from ${OVERLAY_LOC}..."
+        cp -rp "${OVERLAY_LOC}/${OVERLAY_FILE}.gz" ./hpc -n
+        echo "Unzipping overlay ./hpc/${OVERLAY_FILE}.gz..."
+        gunzip "./hpc/${OVERLAY_FILE}.gz" -n
+    fi
+
+    # Launch singularity for the first time
+    echo 'Launching singularity image in WRITE (edit) mode...'
+
+    # Welcome message
+    echo "Run the following to initialize ${PROJECT}:"
+    echo "  (1) create conda environment: 'conda env create -f environment.yml'"
+    echo "  (2) activate conda environment: 'conda activate gpudrive'"
+    echo "  (3) install gpudrive: 'poetry install'"
+
+    # Launch singularity image in write mode
+    singularity exec --nv --overlay "./hpc/${OVERLAY_FILE}:rw" \
+        "${SINGULARITY_IMAGE}" \
+        /bin/bash
+
+else  # Overlay Singularity image and overlay file exist
+
+    # Launch singularity
+    echo 'Launching singularity image in OVERLAY (use) mode...'
+
+    # Welcome message
+    echo "Run the following to activate the Python environment:"
+    echo "  (1) activate conda environment: 'conda activate gpudrive'"
+
+    # Launch singularity image in use mode
+    singularity exec --nv --overlay "./hpc/${OVERLAY_FILE}:ro" \
+        "${SINGULARITY_IMAGE}" \
+        /bin/bash
+
+fi