Fix conda env

.github/workflows/ci.yml

@@ -18,7 +18,7 @@ jobs:
         id: cache-key
         run: |
           # Create a hash from the environment file and PyTorch CPU install string
-          ENV_HASH=$(sha256sum toolbox_env_conda.yml | cut -d' ' -f1)
+          ENV_HASH=$(sha256sum fridata_env_conda.yml | cut -d' ' -f1)
           PYTORCH_SPEC="pip-torch-torchvision-cpu-whl"
           CACHE_KEY="micromamba-${{ runner.os }}-${ENV_HASH}-${PYTORCH_SPEC}"
           echo "key=${CACHE_KEY}" >> $GITHUB_OUTPUT
@@ -36,8 +36,8 @@ jobs:
         if: steps.cache-restore.outputs.cache-hit != 'true'
         uses: mamba-org/setup-micromamba@v2
         with:
-          environment-file: toolbox_env_conda.yml
-          environment-name: tbe
+          environment-file: fridata_env_conda.yml
+          environment-name: fridata_env
           cache-environment: false
           cache-downloads: false
           micromamba-binary-path: ~/micromamba-bin/micromamba
@@ -59,9 +59,9 @@ jobs:
           # CPU wheels from PyTorch are the usual CI approach. Install torchvision from the
           # same index as torch so C++/dispatch ops (e.g. torchvision::nms) match conda-forge
           # transformers alone pulls torchvision that can mismatch a pip-only torch build.
-          ~/micromamba-bin/micromamba run -n tbe python -m pip install --upgrade pip
-          ~/micromamba-bin/micromamba run -n tbe python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
-          ~/micromamba-bin/micromamba run -n tbe python -m pip install esm
+          ~/micromamba-bin/micromamba run -n fridata_env python -m pip install --upgrade pip
+          ~/micromamba-bin/micromamba run -n fridata_env python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+          ~/micromamba-bin/micromamba run -n fridata_env python -m pip install esm
 
       - name: Save micromamba environment cache
         if: steps.cache-restore.outputs.cache-hit != 'true'
@@ -74,4 +74,4 @@ jobs:
 
       - name: Run tests
         run: |
-          ~/micromamba-bin/micromamba run -n tbe pytest ./tests
+          ~/micromamba-bin/micromamba run -n fridata_env pytest ./tests

README.md

@@ -34,20 +34,26 @@ conda install -n base mamba
 4. Create a mamba environment
 
 ```
-mamba create -f toolbox_env_conda.yml
+mamba create -f fridata_env_conda.yml
 ```
 
-5. Activate mamba shell hook
+5. Install pytorch using dedicated script.
+
+```
+./scripts/install_pytorch.sh
+```
+
+6. Activate mamba shell hook
 
 ```
 # Choose your shell type. Could be one of these: {bash,cmd.exe,dash,fish,nu,posix,powershell,tcsh,xonsh,zsh}
 eval "$(mamba shell hook --shell <replace with shell type>)"
 ```
 
-6. Activate the mamba environment
+7. Activate the mamba environment
 
 ```
-mamba activate tbe
+mamba activate fridata_env
 ```
 
 ## Running tests

demo.ipynb

@@ -57,7 +57,7 @@
     "export AFDB_PATH=\"$PATH_TO_DATA/afdb/structures\"\n",
     "\n",
     "eval \"$(conda shell.bash hook)\"\n",
-    "conda activate tbe\n",
+    "conda activate fridata_env\n",
     "\n",
     "fridata generate_data \\\n",
     " -t dataset \\\n",
@@ -92,7 +92,7 @@
     "export PATH_TO_DATASET_JSON=\"$PATH_TO_DATA/datasets/AFDB-subset--testset_inference_over/dataset.json\"\n",
     "\n",
     "eval \"$(conda shell.bash hook)\"\n",
-    "conda activate tbe\n",
+    "conda activate fridata_env\n",
     "\n",
     "fridata create_dashboard \\\n",
     " --dataset ${PATH_TO_DATASET_JSON}"
@@ -161,7 +161,7 @@
     "export PATH_TO_DATA=\"<path to data>\"\n",
     "\n",
     "eval \"$(conda shell.bash hook)\"\n",
-    "conda activate tbe\n",
+    "conda activate fridata_env\n",
     "\n",
     "fridata load --file-path \"$PATH_TO_DATA/datasets/AFDB-subset--testset_inference_over\""
    ]
@@ -208,7 +208,7 @@
     "export PATH_TO_COORDINATES=\"$PATH_TO_DATA/coordinates/AFDB-subset--_test_dask/batch_0_ca.h5\"\n",
     "\n",
     "eval \"$(conda shell.bash hook)\"\n",
-    "conda activate tbe\n",
+    "conda activate fridata_env\n",
     "\n",
     "fridata inspect_h5 --help"
    ]

toolbox_env_conda.yml → fridata_env_conda.yml

@@ -1,4 +1,4 @@
-name: tbe
+name: fridata_env
 channels:
   - conda-forge
   - defaults
@@ -26,8 +26,6 @@ dependencies:
   - biopython
   - bokeh>=2.4.2,<3.0.0
   - pdbfixer
-  # PyTorch and ESM installed separately via install.sh (GPU or CPU variant)
-  # ESM requires PyTorch to be installed first
   - mamba
   - transformers
   - tqdm

scripts/hpc/initialize_slurm.sh

@@ -58,7 +58,7 @@ fi
 conda config --add pkgs_dirs "$CONDA_DIR"
 
 # Create environment from base YAML (without PyTorch)
-conda env create --prefix $CONDA_ENV_PATH --file "$DEEPFRI_PATH/FRIdata/toolbox_env_conda.yml"
+conda env create --prefix $CONDA_ENV_PATH --file "$DEEPFRI_PATH/FRIdata/fridata_env_conda.yml"
 
 conda config --set auto_activate_base false
 

scripts/install_pytorch.sh

@@ -0,0 +1,146 @@
+#!/bin/bash
+# Install a PyTorch build matching the machine's CUDA driver, then install ESM.
+#
+# Fixes errors such as:
+#   UserWarning: CUDA initialization: The NVIDIA driver on your system is too
+#   old (found version 12040). ...
+# which happen when the installed PyTorch was compiled against a newer CUDA
+# toolkit than the installed NVIDIA driver supports.
+#
+# The script reads the maximum CUDA version the driver supports from
+# `nvidia-smi`, selects the highest PyTorch wheel (cuXXX) that is not newer
+# than that, and installs it into the *currently active* environment via pip.
+# If no NVIDIA GPU/driver is found it installs the CPU-only build.
+#
+# Usage: ./install_pytorch.sh [--cpu] [--dry-run]
+#   --cpu      Force the CPU-only PyTorch build (skip GPU detection)
+#   --dry-run  Print what would be installed without installing anything
+
+set -euo pipefail
+
+# PyTorch wheel CUDA versions, ascending. Update as new wheels are released.
+# https://pytorch.org/get-started/locally/ and https://download.pytorch.org/whl/
+SUPPORTED_CUDA=(11.8 12.1 12.4 12.6 12.8)
+
+CPU_ONLY=false
+DRY_RUN=false
+for arg in "$@"; do
+    case "$arg" in
+        --cpu)     CPU_ONLY=true ;;
+        --dry-run) DRY_RUN=true ;;
+        -h|--help)
+            grep '^#' "$0" | sed 's/^# \{0,1\}//'
+            exit 0
+            ;;
+        *)
+            echo "Error: unknown argument '$arg'" >&2
+            echo "Usage: ./install_pytorch.sh [--cpu] [--dry-run]" >&2
+            exit 1
+            ;;
+    esac
+done
+
+# --- Verify an environment is active ------------------------------------------
+if ! command -v python >/dev/null 2>&1; then
+    echo "Error: no 'python' found on PATH. Activate your conda/venv first." >&2
+    exit 1
+fi
+
+if [ -n "${CONDA_PREFIX:-}" ]; then
+    echo "Target environment: $CONDA_PREFIX"
+    if [ "$(basename "$CONDA_PREFIX")" = "base" ]; then
+        echo "Warning: the conda 'base' environment is active. It is strongly" >&2
+        echo "         recommended to activate a dedicated environment first." >&2
+    fi
+else
+    echo "Warning: CONDA_PREFIX is not set; installing into '$(command -v python)'." >&2
+fi
+
+# Turn a X.Y version into a comparable integer, e.g. 12.4 -> 1204.
+cuda_to_int() {
+    local major minor
+    major="${1%%.*}"
+    minor="${1#*.}"
+    printf '%d' "$((major * 100 + minor))"
+}
+
+# --- Detect the driver's maximum supported CUDA version -----------------------
+CUDA_TAG=""
+if [ "$CPU_ONLY" = true ]; then
+    echo "Forcing CPU-only PyTorch (--cpu)."
+elif ! command -v nvidia-smi >/dev/null 2>&1; then
+    echo "No 'nvidia-smi' found; assuming no NVIDIA GPU -> CPU-only PyTorch."
+    CPU_ONLY=true
+else
+    # The 'CUDA Version' field in nvidia-smi reports the highest CUDA runtime
+    # the installed driver can support.
+    DRIVER_CUDA="$(nvidia-smi 2>/dev/null \
+        | grep -oE 'CUDA Version: [0-9]+\.[0-9]+' \
+        | grep -oE '[0-9]+\.[0-9]+' \
+        | head -n1 || true)"
+
+    if [ -z "$DRIVER_CUDA" ]; then
+        echo "Could not read a CUDA version from nvidia-smi; falling back to CPU-only." >&2
+        CPU_ONLY=true
+    else
+        echo "Driver supports up to CUDA $DRIVER_CUDA."
+        driver_int="$(cuda_to_int "$DRIVER_CUDA")"
+
+        # Pick the highest supported wheel that is <= the driver's max CUDA.
+        chosen=""
+        for v in "${SUPPORTED_CUDA[@]}"; do
+            if [ "$(cuda_to_int "$v")" -le "$driver_int" ]; then
+                chosen="$v"
+            fi
+        done
+
+        if [ -z "$chosen" ]; then
+            echo "Driver CUDA $DRIVER_CUDA is older than the oldest available" >&2
+            echo "PyTorch CUDA wheel (${SUPPORTED_CUDA[0]}); falling back to CPU-only." >&2
+            CPU_ONLY=true
+        else
+            CUDA_TAG="cu${chosen//./}"
+            echo "Selected PyTorch CUDA build: $CUDA_TAG (for CUDA $chosen)."
+        fi
+    fi
+fi
+
+# --- Build and run the install command ----------------------------------------
+if [ "$CPU_ONLY" = true ]; then
+    INDEX_URL="https://download.pytorch.org/whl/cpu"
+else
+    INDEX_URL="https://download.pytorch.org/whl/${CUDA_TAG}"
+fi
+
+PIP_CMD=(python -m pip install torch torchvision torchaudio --index-url "$INDEX_URL")
+
+echo
+echo "Will run: ${PIP_CMD[*]}"
+echo "Then:     python -m pip install esm"
+
+if [ "$DRY_RUN" = true ]; then
+    echo "(--dry-run: nothing installed)"
+    exit 0
+fi
+
+echo
+echo "Installing PyTorch..."
+"${PIP_CMD[@]}"
+
+# ESM requires PyTorch to be installed first.
+echo "Installing ESM..."
+python -m pip install esm
+
+# --- Verify the installation --------------------------------------------------
+echo "Verifying PyTorch installation..."
+python - <<'PY'
+import torch
+print(f"torch {torch.__version__}")
+if torch.cuda.is_available():
+    print(f"CUDA available: True (torch built for CUDA {torch.version.cuda})")
+    print(f"Detected GPU: {torch.cuda.get_device_name(0)}")
+else:
+    print("CUDA available: False (CPU-only build or no usable GPU)")
+PY
+
+echo "Installation complete!"