File size: 4,701 Bytes
6ae7efa e40c47f 6ae7efa 5bdf411 6ae7efa e40c47f 6ae7efa e40c47f 6ae7efa e40c47f 6ae7efa e40c47f 6ae7efa 5bdf411 6ae7efa e40c47f 6ae7efa e40c47f 6ae7efa e40c47f 6ae7efa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | #!/usr/bin/env bash
# Usage: bash setup_env.sh [environment-prefix] [all|both|core|benchmark|evo2]
set -euo pipefail
# Keep unrelated per-user Python packages out of dependency resolution/imports.
export PYTHONNOUSERSITE=1
unset PYTHONPATH
project_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
environment_name="${1:-rnaseek-release}"
profile="${2:-all}"
if [[ "$profile" != "all" && "$profile" != "both" && "$profile" != "core" && "$profile" != "benchmark" && "$profile" != "evo2" ]]; then
echo "Profile must be all, both (core + benchmark), core, benchmark, or evo2." >&2
exit 2
fi
if [[ "$environment_name" == "base" || ! "$environment_name" =~ ^[A-Za-z0-9_-]+$ ]]; then
echo "Choose a new environment name containing only letters, digits, underscores or hyphens." >&2
exit 2
fi
if ! command -v conda >/dev/null 2>&1; then
echo "Install Miniforge or Miniconda, reopen your terminal, then rerun this script." >&2
exit 2
fi
# A pre-existing environment may contain unrelated work. Never update it here.
conda_base="$(conda info --base)"
environment_list="$(conda env list --json)"
profiles=()
[[ "$profile" == "all" || "$profile" == "both" || "$profile" == "core" ]] && profiles+=(core)
[[ "$profile" == "all" || "$profile" == "both" || "$profile" == "benchmark" ]] && profiles+=(benchmark)
[[ "$profile" == "all" || "$profile" == "evo2" ]] && profiles+=(evo2)
if [[ "$profile" == "all" || "$profile" == "evo2" ]]; then
"$conda_base/bin/python" -c 'import platform,sys; sys.exit(0 if sys.platform == "linux" and platform.machine() == "x86_64" else "The pinned Evo 2 GPU backend requires Linux x86-64 (WSL2 is an option on Windows). Use profile both for core/encoder environments only.")'
fi
# Check all requested names before creating either environment.
for selected in "${profiles[@]}"; do
target_name="$environment_name"
[[ "$selected" == "benchmark" ]] && target_name="$environment_name-benchmark"
[[ "$selected" == "evo2" ]] && target_name="$environment_name-evo2"
environment_exists="$("$conda_base/bin/python" -c \
'import json,sys,pathlib; name=sys.argv[1]; print(int(any(pathlib.Path(p).name == name for p in json.load(sys.stdin)["envs"])))' \
"$target_name" <<< "$environment_list")"
if [[ "$environment_exists" == "1" ]]; then
echo "Requested environment already exists. Supply a new prefix or choose the missing profile." >&2
exit 2
fi
done
for selected in "${profiles[@]}"; do
target_name="$environment_name"
requirements="$project_dir/requirements.txt"
constraints="$project_dir/constraints-core.txt"
if [[ "$selected" == "benchmark" ]]; then
target_name="$environment_name-benchmark"
requirements="$project_dir/requirements-benchmark.txt"
constraints="$project_dir/constraints-benchmark.txt"
elif [[ "$selected" == "evo2" ]]; then
target_name="$environment_name-evo2"
requirements="$project_dir/requirements-evo2.txt"
constraints="$project_dir/constraints-evo2.txt"
fi
conda create --yes --name "$target_name" --override-channels --channel conda-forge python=3.11 pip
conda env config vars set --name "$target_name" PYTHONNOUSERSITE=1
conda run --no-capture-output --name "$target_name" python -m pip install -r "$requirements" -c "$constraints"
if [[ "$selected" == "evo2" ]]; then
# Official prebuilt wheel matching Python 3.11 / Torch 2.7 / CUDA 12 / CXX11 ABI.
# Install after Torch; avoid a local CUDA-toolkit source-build requirement.
conda run --no-capture-output --name "$target_name" python -m pip install -c "$constraints" \
'https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiTRUE-cp311-cp311-linux_x86_64.whl'
fi
conda run --no-capture-output --name "$target_name" python -m pip check
if [[ "$selected" == "core" ]]; then
conda run --no-capture-output --name "$target_name" python -c \
'import torch,transformers,trl,nbformat,nbclient,RNA; print("Core environment ready; GPUs:",torch.cuda.device_count())'
elif [[ "$selected" == "benchmark" ]]; then
conda run --no-capture-output --name "$target_name" python -c \
'import torch,transformers,multimolecule; print("Benchmark environment ready; GPUs:",torch.cuda.device_count())'
else
conda run --no-capture-output --name "$target_name" python -c \
'import torch,evo2; from vortex.model.model import get_block; print("Evo 2 environment ready; GPUs:",torch.cuda.device_count())'
fi
printf '\nActivate with: conda activate %s\n' "$target_name"
done
|