diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 307986d..94aa5a2 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/README.md b/README.md index 3ad5112..8f1263a 100644 --- a/README.md +++ b/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 )" ``` -6. Activate the mamba environment +7. Activate the mamba environment ``` -mamba activate tbe +mamba activate fridata_env ``` ## Running tests diff --git a/demo.ipynb b/demo.ipynb index 38b6537..4f64507 100644 --- a/demo.ipynb +++ b/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=\"\"\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" ] diff --git a/toolbox_env_conda.yml b/fridata_env_conda.yml similarity index 79% rename from toolbox_env_conda.yml rename to fridata_env_conda.yml index 5973380..a96813f 100644 --- a/toolbox_env_conda.yml +++ b/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 diff --git a/scripts/hpc/initialize_slurm.sh b/scripts/hpc/initialize_slurm.sh index 96c733d..cfe4452 100644 --- a/scripts/hpc/initialize_slurm.sh +++ b/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 diff --git a/scripts/install_pytorch.sh b/scripts/install_pytorch.sh new file mode 100755 index 0000000..3ae3f32 --- /dev/null +++ b/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!"