setup_env.sh

#!/bin/bash

# ==============================================================================
# NeuroMamba environment installer
# Sets up a Python venv with PyTorch + Lightning + DeepSpeed + Mamba SSM.
# ==============================================================================

# [User configuration]
# ==============================================================================
# Virtual-environment path (absolute path recommended)
VENV_PATH="${NEUROMAMBA_VENV_PATH:-$HOME/neuromamba_env}"

# Common pip packages (DeepSpeed installed separately below).
# Note: pytorch_lightning==1.9.5 may have compatibility issues with newer torch
# versions, but is kept here for reproducibility with the paper.
COMMON_PACKAGES="timm einops wandb pytorch_lightning==1.9.5 monai nilearn nibabel scikit-learn mup psmpy h5py tensorboard mpi4py mpich "

# Target CUDA architecture for DeepSpeed compilation (B200 = 10.0)
TARGET_CUDA_ARCH="10.0"

# ==============================================================================
# [1] Create the virtual environment (inheriting system site-packages)
# ==============================================================================
echo "========================================================"
echo "[INFO] Creating virtual environment."
echo "       Path:    ${VENV_PATH}"
echo "       Option:  --system-site-packages (inherit system libraries)"
echo "========================================================"

if [ ! -d "$(dirname "$VENV_PATH")" ]; then
    echo "[INFO] Creating parent directory: $(dirname "$VENV_PATH")"
    mkdir -p "$(dirname "$VENV_PATH")"
fi

if [ -d "$VENV_PATH" ]; then
    echo "[WARN] Path already exists: ${VENV_PATH}"
    echo "[INFO] Reusing existing environment; only adding missing packages."
else
    python3 -m venv --system-site-packages "$VENV_PATH"
    if [ $? -ne 0 ]; then
        echo "[ERROR] Failed to create virtual environment."
        exit 1
    fi
    echo "[INFO] Virtual environment created."
fi

# ==============================================================================
# [2] Activate environment and update pip
# ==============================================================================
echo "[INFO] Activating virtual environment..."
source "${VENV_PATH}/bin/activate"

echo "[INFO] Upgrading pip..."
pip install --upgrade pip

# ==============================================================================
# [3] Install common packages
# ==============================================================================
echo "========================================================"
echo "[INFO] Installing common packages..."
echo "       Packages: ${COMMON_PACKAGES}"
echo "========================================================"

pip install ${COMMON_PACKAGES}

# ==============================================================================
# [4] Compile and install DeepSpeed with target CUDA arch
# ==============================================================================
echo "========================================================"
echo "[INFO] Installing DeepSpeed..."
echo "       Target arch: ${TARGET_CUDA_ARCH}"
echo "       Ops: Fused Adam, CPU Adam"
echo "========================================================"

export TORCH_CUDA_ARCH_LIST="${TARGET_CUDA_ARCH}"
export DS_BUILD_FUSED_ADAM=1
export DS_BUILD_CPU_ADAM=1

# --no-build-isolation is required so the build sees the system-installed
# PyTorch and compiles ops against it.
pip install deepspeed --no-cache-dir --no-build-isolation

if [ $? -ne 0 ]; then
    echo "[ERROR] DeepSpeed build/install failed."
    echo "       Verify CUDA toolkit and PyTorch version compatibility."
    exit 1
fi

# ==============================================================================
# [4.5] Install mamba-ssm + causal-conv1d (NeuroMamba's core dependencies)
# ==============================================================================
echo "========================================================"
echo "[INFO] Installing mamba-ssm and causal-conv1d..."
echo "       Target arch: ${TARGET_CUDA_ARCH}"
echo "========================================================"

# --no-build-isolation uses the current env's build tools — ensure they exist.
# setuptools is pinned < 80: newer versions drop pkg_resources, which
# pytorch_lightning 1.9.5 (lightning_fabric) still imports at startup.
pip install --upgrade "setuptools>=70.1,<80" wheel ninja

# Install causal-conv1d first (mamba-ssm depends on it)
pip install causal-conv1d --no-cache-dir --no-build-isolation

if [ $? -ne 0 ]; then
    echo "[ERROR] causal-conv1d build failed."
    exit 1
fi

# Install mamba-ssm
pip install mamba-ssm --no-cache-dir --no-build-isolation

if [ $? -ne 0 ]; then
    echo "[ERROR] mamba-ssm build failed."
    exit 1
fi

echo "[INFO] mamba-ssm and causal-conv1d installed."

# ==============================================================================
# [5] Verify and finish
# ==============================================================================
echo "========================================================"
echo "[RESULT] Installation complete."
echo "========================================================"

echo "1) Python path:"
which python
echo ""

echo "2) Installed package versions:"
pip list | grep -E "torch|deepspeed|monai|lightning|mamba|causal"
echo ""

echo "3) mamba-ssm import test:"
python -c "from mamba_ssm import Mamba, Mamba2; print('  mamba-ssm OK')"
if [ $? -ne 0 ]; then
    echo "[WARN] mamba-ssm import failed. Manual verification required."
fi

echo "--------------------------------------------------------"
echo "Activate the environment with:"
echo "source ${VENV_PATH}/bin/activate"
echo ""
echo "NeuroMamba test:"
echo "cd \${PROJECT_ROOT}"  # path to your local NeuroMamba clone
echo "python -c \"from project.module.models.fmamba import FMamba; print('FMamba OK')\""