Add python libraries in a separate conda env (#333)

This change adds back the python libraries removed in #322. Rather than adding them to `CONDA_PACKAGES_2`, we create a third conda environment explicitly for python libraries. This keeps the environments more distinct: one for CLI tools, another for samtools (separated due to dependency version conflicts), and a third for python libraries. The python libraries environment is intended to be the user's main python environment, so we should not need to manipulate `PYTHONPATH`. This python libraries environment can either be an existing environment or a new one.

PHP-132843

Author: pantherman594

features/src/workbench-tools/devcontainer-feature.json

@@ -4,6 +4,16 @@
   "name": "Workbench Tools",
   "description": "Installs common tools for Workbench Apps. Currently it only supports Debian-based systems (e.g. Ubuntu) on x86_64.",
   "options": {
+      "libEnv": {
+          "type": "string",
+          "default": "/opt/conda/envs/workbench-ds",
+          "description": "The path to the conda environment where Python libraries (e.g. numpy, plotly, scipy) should be installed. This could be a pre-existing environment or a new one."
+      },
+      "libPythonVersion": {
+          "type": "string",
+          "default": "3.14",
+          "description": "The Python version to use for the installed libraries."
+      },
       "cloud": {
           "type": "string",
           "default": "",

features/src/workbench-tools/install-conda.sh

@@ -19,7 +19,6 @@ install_miniforge() {
     local download_url
     download_url="https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-$(uname -m).sh"
 
-    check_packages curl ca-certificates
     mkdir -p /tmp/miniforge
     (
         cd /tmp/miniforge

features/src/workbench-tools/install.sh

@@ -10,6 +10,8 @@ set -o xtrace
 
 readonly CLOUD="${CLOUD:-""}"
 readonly USERNAME="${USERNAME:-"root"}"
+readonly LIBRARIES_ENV_DIR="${LIBENV:-"/opt/conda/envs/workbench-ds"}"
+readonly LIB_PYTHON_VERSION="${LIBPYTHONVERSION:-"3.14"}"
 USER_HOME_DIR="${USERHOMEDIR:-"/home/${USERNAME}"}"
 if [[ "${USER_HOME_DIR}" == "/home/root" ]]; then
     USER_HOME_DIR="/root"
@@ -68,66 +70,122 @@ if ! mamba --version &>/dev/null; then
 fi
 
 # Install the samtools family of tools in a separate environment since some of
-# the other tools depend on old versions of these. This will take priority in
-# the PATH.
-CONDA_PACKAGES_1=(
-    "bcftools"
-    "htslib" # includes bgzip and tabix
-    "samtools"
+# the other tools depend on old versions of these.
+readonly CONDA_PACKAGES_SAMTOOLS=(
+    "bioconda::bcftools>=1.23"
+    "bioconda::htslib>=1.23" # includes bgzip and tabix
+    "bioconda::samtools>=1.23"
 )
-
-CONDA_PACKAGES_2=(
-    "python=3.9"
-    "pip"
-    "perl==5.32.1"
-    "bedtools"
+readonly SAMTOOLS_ENV_DIR="${WORKBENCH_TOOLS_DIR}/samtools"
+
+# Environment 2 contains the genomics CLI tools. They will be added to the
+# PATH but will not be usable as Python libraries.
+readonly CONDA_PACKAGES_BINARIES=(
+    "conda-forge::python"
+    "conda-forge::pip"
+    "conda-forge::perl>=5.32"
+    "bioconda::bedtools"
     "conda-forge::bgenix"
     "conda-forge::cromwell"
-    "ensembl-vep>=115.1"
-    "nextflow"
-    "plink"
-    "plink2"
-    "regenie"
-    "vcftools"
+    "bioconda::ensembl-vep>=115"
+    "bioconda::nextflow"
+    "bioconda::plink"
+    "bioconda::plink2"
+    "bioconda::regenie"
+    "bioconda::vcftools"
+)
+readonly BINARIES_ENV_DIR="${WORKBENCH_TOOLS_DIR}/binaries"
+
+# Environment 3 contains data science Python libraries. These should be
+# accessible from the user's default Python environment, which is why we install
+# them separately and give the user control over whether to inject them into an
+# existing environment or create a new one.
+CONDA_PACKAGES_LIBRARIES=(
+    "conda-forge::google-cloud-storage"
+    "conda-forge::ipykernel"
+    "conda-forge::ipywidgets"
+    "conda-forge::jupyter"
+    "conda-forge::openai"
+    "conda-forge::matplotlib"
+    "conda-forge::numpy"
+    "conda-forge::plotly"
+    "conda-forge::pandas"
+    "conda-forge::seaborn"
+    "conda-forge::scikit-learn"
+    "conda-forge::scipy"
+    "conda-forge::tqdm"
 )
 
+# Build isolated environments
 mkdir -p "${WORKBENCH_TOOLS_DIR}"
-mamba create --prefix "${WORKBENCH_TOOLS_DIR}/1" -c bioconda -y "${CONDA_PACKAGES_1[@]}"
-mamba create --prefix "${WORKBENCH_TOOLS_DIR}/2" -c bioconda -y "${CONDA_PACKAGES_2[@]}"
+echo "Building Environment 1 (Samtools family)..."
+mamba create --prefix "${SAMTOOLS_ENV_DIR}" -y "${CONDA_PACKAGES_SAMTOOLS[@]}"
+
+echo "Building Environment 2 (Genomics CLI Tools)..."
+mamba create --prefix "${BINARIES_ENV_DIR}" -y "${CONDA_PACKAGES_BINARIES[@]}"
+
+echo "Building Environment 3 (Python Libraries)..."
+LIB_ENV_EXISTS=0
+
+if [ -d "${LIBRARIES_ENV_DIR}" ]; then
+    # SCENARIO A: Target environment already exists on host. Inject packages into it.
+    LIB_ENV_EXISTS=1
+    echo "Host environment detected at ${LIBRARIES_ENV_DIR}. Injecting data science packages..."
+
+    if mamba list -p "${LIBRARIES_ENV_DIR}" --full-name python --json | jq -e 'length == 0' >/dev/null; then
+        echo "No Python installation found in host environment. Adding python=${LIB_PYTHON_VERSION} to package list."
+        CONDA_PACKAGES_LIBRARIES+=("conda-forge::python=${LIB_PYTHON_VERSION}")
+    fi
+    mamba install --prefix "${LIBRARIES_ENV_DIR}" -y "${CONDA_PACKAGES_LIBRARIES[@]}"
+else
+    # SCENARIO B: Target environment does not exist. Create it from scratch.
+    echo "No host environment found. Creating standalone environment at ${LIBRARIES_ENV_DIR}..."
+    mkdir -p "$(dirname "${LIBRARIES_ENV_DIR}")"
+
+    CONDA_PACKAGES_LIBRARIES+=("conda-forge::python=${LIB_PYTHON_VERSION}")
+    mamba create --prefix "${LIBRARIES_ENV_DIR}" -y "${CONDA_PACKAGES_LIBRARIES[@]}"
+fi
 
 # Install dsub via pip if on GCP. The conda version is outdated.
+# dsub is installed in LIBRARIES_ENV_DIR because it can be used as a Python
+# library, and users may want to install additional packages alongside it.
+# PYTHONNOUSERSITE=1 prevents pip from seeing/modifying packages in user site-packages.
 if [[ "${CLOUD}" == "gcp" ]]; then
-    "${WORKBENCH_TOOLS_DIR}/2/bin/pip" install dsub
+    PYTHONNOUSERSITE=1 "${LIBRARIES_ENV_DIR}/bin/pip" install dsub
 fi
 
 # Force the perl and python scripts to use the correct perl/python
-find -L "${WORKBENCH_TOOLS_DIR}/2/bin" -type f -executable -exec \
+find -L "${BINARIES_ENV_DIR}/bin" -type f -executable -exec \
     sed -i --follow-symlinks \
-        -e "1s|^#\!/usr/bin/env perl\\r\?$|#\!${WORKBENCH_TOOLS_DIR}/2/bin/perl|" \
-        -e "1s|^#\!/usr/bin/env python\\r\?$|#\!${WORKBENCH_TOOLS_DIR}/2/bin/python|" {} \;
+        -e "1s|^#\!/usr/bin/env perl\\r\?$|#\!${BINARIES_ENV_DIR}/bin/perl|" \
+        -e "1s|^#\!/usr/bin/env python\\r\?$|#\!${BINARIES_ENV_DIR}/bin/python|" {} \;
 
-# Make the login user the owner of the conda environment
+# Make the login user the owner of the conda environments
 chown -R "${USERNAME}:" "${WORKBENCH_TOOLS_DIR}"
+chown -R "${USERNAME}:" "${LIBRARIES_ENV_DIR}"
 
 {
+    echo "# Workbench Tools Configuration"
+    
+    # If we created a standalone Python libraries environment from scratch, make it the default terminal Python.
+    # If it already existed (LIB_ENV_EXISTS=1), we leave the host image's PATH untouched to prevent shadowing.
+    if [[ "${LIB_ENV_EXISTS}" == "0" ]]; then
+        # shellcheck disable=SC2016 # we want $PATH to be evaluated at runtime
+        printf 'export PATH="%s:$PATH"\n' "${LIBRARIES_ENV_DIR}/bin"
+    fi
+    
     # Set PATH to include workbench-tools binaries
     # shellcheck disable=SC2016 # we want $PATH to be evaluated at runtime
-    printf 'export PATH="%s:$PATH"\n' "${WORKBENCH_TOOLS_DIR}/1/bin:${WORKBENCH_TOOLS_DIR}/2/bin"
+    printf 'export PATH="$PATH:%s"\n' "${SAMTOOLS_ENV_DIR}/bin:${BINARIES_ENV_DIR}/bin"
 
     # Set CROMWELL_JAR environment variable
-    printf 'export CROMWELL_JAR="%s"\n' "${WORKBENCH_TOOLS_DIR}/2/share/cromwell/cromwell.jar"
-
-    # Make dsub a function that includes the correct PYTHONPATH. NeMo sets
-    # PYTHONPATH so we need to override it here. We use a function instead of an
-    # alias because aliases are not expanded in non-interactive shells.
-    # shellcheck disable=SC2016 # we want $PYTHONPATH to be evaluated at runtime
-    printf 'function dsub() (PYTHONPATH="%s/2/lib/python3.9/site-packages:${PYTHONPATH:-}" "%s/2/bin/dsub" "$@")\n' "${WORKBENCH_TOOLS_DIR}" "${WORKBENCH_TOOLS_DIR}" 
+    printf 'export CROMWELL_JAR="%s"\n' "${BINARIES_ENV_DIR}/share/cromwell/cromwell.jar"
 } >> "${USER_HOME_DIR}/.bashrc"
 
 # Allow .bashrc to be sourced in non-interactive shells
 sed -i '/^# If not running interactively/,/esac/d' "${USER_HOME_DIR}/.bashrc" || true
 
 # Make sure the login user is the owner of their .bashrc
-chown -R "${USERNAME}:" "${USER_HOME_DIR}/.bashrc"
+chown "${USERNAME}:" "${USER_HOME_DIR}/.bashrc"
 
-echo "Done!"
+echo "Workbench tools installation complete!"

src/custom-workbench-jupyter-template/.devcontainer.json

@@ -22,6 +22,7 @@
   ],
   "features": {
     "./.devcontainer/features/workbench-tools": {
+      "libEnv": "/opt/conda/envs/jupyter", // Use the jupyter conda environment
       "cloud": "${templateOption:cloud}",
       "username": "jupyter",
       "userHomeDir": "/home/jupyter"

src/jupyter-aou/.devcontainer.json

@@ -11,6 +11,7 @@
   "postStartCommand": "./startupscript/remount-on-restart.sh jupyter /home/jupyter \"${templateOption:cloud}\" \"${templateOption:login}\"",
   "features": {
     "./.devcontainer/features/workbench-tools": {
+      "libEnv": "/opt/conda/envs/jupyter", // Use the jupyter conda environment
       "cloud": "${templateOption:cloud}",
       "username": "jupyter",
       "userHomeDir": "/home/jupyter"

src/nemo_jupyter/.devcontainer.json

@@ -25,6 +25,7 @@
     },
     "ghcr.io/dhoeric/features/google-cloud-cli@sha256:fa5d894718825c5ad8009ac8f2c9f0cea3d1661eb108a9d465cba9f3fc48965f": {},
     "./.devcontainer/features/workbench-tools": {
+      "libPythonVersion": "3.12", // Must match python version in nemo image
       "cloud": "${templateOption:cloud}",
       "username": "jupyter",
       "userHomeDir": "/home/jupyter"

src/nemo_jupyter/Dockerfile

@@ -17,9 +17,15 @@ RUN useradd --uid ${NB_UID} --gid ${NB_GID} --create-home --home-dir ${WORKDIR}
     # Fix ownership for common dirs
     && mkdir -p /workspace \
     && chown -R ${NB_UID}:${NB_GID} ${WORKDIR} /workspace /tmp \
-    && chown -R ${NB_UID}:${NB_GID} /opt/conda || true \
-    # Add alias for weightsbiases
-    && printf 'alias weightsbiases="/usr/local/bin/wb"\nalias wb="/usr/bin/wb"\n' >> ${WORKDIR}/.bashrc
+    && chown -R ${NB_UID}:${NB_GID} /opt/conda || true
+
+# Add shell functions to override wb commands (functions work in non-interactive shells, unlike aliases)
+RUN <<EOF cat >> ${WORKDIR}/.bashrc
+weightsbiases() { /usr/local/bin/wb "\$@"; }
+wb() { /usr/bin/wb "\$@"; }
+wandb() { /opt/venv/bin/wb "\$@"; }
+export -f weightsbiases wb wandb
+EOF
 
 # Environment and working directory
 ENV HOME=${WORKDIR}

src/nemo_jupyter_aou/.devcontainer.json

@@ -26,6 +26,7 @@
     },
     "ghcr.io/dhoeric/features/google-cloud-cli@sha256:fa5d894718825c5ad8009ac8f2c9f0cea3d1661eb108a9d465cba9f3fc48965f": {},
     "./.devcontainer/features/workbench-tools": {
+      "libPythonVersion": "3.12", // Must match python version in nemo image
       "cloud": "${templateOption:cloud}",
       "username": "jupyter",
       "userHomeDir": "/home/jupyter"

startupscript/post-startup.sh

@@ -201,7 +201,7 @@ cat << EOF >> "${USER_BASHRC}"
 ### BEGIN: Workbench-specific customizations ###
 
 # Prepend "/usr/bin" (if not already in the path)
-if [[ "\${PATH}:" != "/usr/bin:"* ]]; then
+if [[ ":\${PATH}:" != *":/usr/bin:"* ]]; then
   export PATH=/usr/bin:\${PATH}
 fi
 

test/test.sh

@@ -11,16 +11,18 @@ function check() {
     check_user "${TEST_USER}" "$@"
 }
 
+sudo -u "${TEST_USER}" bash -l -c "echo PATH: $PATH"
+
 # Template specific tests
 check "gcsfuse" gcsfuse -v
 check "wb cli" wb version
 check "fuse.conf user_allow_other" grep -qE "^[[:space:]]*[^#]*user_allow_other" "/etc/fuse.conf"
 
 # The workbench-tools feature should install these
 if [[ "$HAS_WORKBENCH_TOOLS" == "true" ]]; then
-    check "python3" python3 --version
+    check "python3" "which python3 && python3 --version"
     check "python3: venv" 'python3 -c "import venv"'
-    check "pip3" pip3 --version
+    check "pip3" "which pip3 && pip3 --version"
     if [[ "$TEMPLATE_ID" != "nemo_jupyter" ]] && [[ "$TEMPLATE_ID" != "nemo_jupyter_aou" ]]; then
         check "cromwell" cromwell --version
     fi
@@ -44,6 +46,21 @@ if [[ "$HAS_WORKBENCH_TOOLS" == "true" ]]; then
     check "vep: filter_vep" "filter_vep --help > /dev/null"
     check "vep: variant_recoder" "variant_recoder --help | head -n10"
     check "vep: haplo" "haplo --help | head -n10"
+
+    # Python packages
+    check "python: google-cloud-storage" 'python3 -c "import google.cloud.storage"'
+    check "python: ipykernel" 'python3 -c "import ipykernel"'
+    check "python: ipywidgets" 'python3 -c "import ipywidgets"'
+    check "python: jupyter" 'python3 -c "import jupyter"'
+    check "python: openai" 'python3 -c "import openai"'
+    check "python: matplotlib" 'python3 -c "import matplotlib"'
+    check "python: numpy" 'python3 -c "import numpy"'
+    check "python: plotly" 'python3 -c "import plotly"'
+    check "python: pandas" 'python3 -c "import pandas"'
+    check "python: seaborn" 'python3 -c "import seaborn"'
+    check "python: scikit-learn" 'python3 -c "import sklearn"'
+    check "python: scipy" 'python3 -c "import scipy"'
+    check "python: tqdm" 'python3 -c "import tqdm"'
 fi
 
 # The postgres-client feature should install these