# Suno GPT Environment Builder

Robust environment setup scripts for Suno GPT with both Conda and Virtualenv support, comprehensive testing, and environment management utilities.

## 📋 Table of Contents

- [Quick Start](#-quick-start)
- [Prerequisites](#-prerequisites)
- [Installation Methods](#-installation-methods)
- [Verification & Testing](#-verification--testing)
- [Environment Management](#-environment-management)
- [Troubleshooting](#-troubleshooting)
- [Package List](#-package-list)
- [File Structure](#-file-structure)
- [Updating Environments](#-updating-environments)
- [Notes](#-notes)
- [Support](#-support)
- [License](#-license)

## 🚀 Quick Start

### Create New Environment with Flash Attention v3 + v2

**Recommended:** This setup installs both FA3 (H100 optimized) and FA2 (compatibility) in one environment.

**Prerequisites:**
- ⚠️⚠️⚠️ WARNING: `glockenspiel/suno_utils` will be updated and install as editable. It might affect the current environment sharing this package if there is any CPython building. *You may want to clone another `glockenspiel` if you want zero side effect*.
- Make sure glockenspiel is on the latest main branch:
```bash
# suggested to clone a new glockenspiel to avoid any side effect
cd ~/projects/glockenspiel
git checkout main && git pull
```
- Know the path to your suno_utils (typically `~/projects/glockenspiel/suno_utils`)

**Usage:**
```bash
./setup_suno_env_fa2_fa3_fixed_commit.sh <env_name> <suno_utils_path>
```

#### Option 1: Tmux Setup (Recommended)

Run the setup script in a tmux session to monitor progress because it will take a long time:

```bash
# Start a tmux session
tmux new -s suno_setup

# Run the setup script with suno_utils path
cd ~/projects/ml_build_env
./setup_suno_env_fa2_fa3_fixed_commit.sh suno_env_auto ~/projects/glockenspiel/suno_utils
```

The script will automatically:
- Validate suno_utils path
- Clone Flash Attention to `/tmp/flash-attention-v3-{env_name}` (commit c5b0c631)
- Clone Flash Attention to `/tmp/flash-attention-{env_name}` for FA2 (commit c5b0c631)
- No interactive prompts!

To detach from tmux: Press `Ctrl+B`, then `D`
To reattach: `tmux attach -t suno_setup`

#### Option 2: Automated Background Setup

Fully automated setup with no interaction required:

```bash
cd ~/projects/ml_build_env
nohup ./setup_suno_env_fa2_fa3_fixed_commit.sh suno_env_auto ~/projects/glockenspiel/suno_utils > setup_suno_env_auto.log 2>&1 &
```

Monitor progress:
```bash
tail -f setup_suno_env_auto.log
```

**Note:**
- Suno utils path is required as second argument
- Flash Attention is cloned fresh to `/tmp` with fixed commit c5b0c631
- No manual Flash Attention repository required!
- Fully non-interactive - runs completely unattended

**Setup time:** ~3-4 hours (FA3 build takes ~2.5 hours on H100)

**What it installs:**
- PyTorch 2.6.0 + CUDA 12.4
- Flash Attention v3 (H100/Hopper optimized)
- Flash Attention v2.8.1 (commit c5b0c631 for compatibility)
- Transformers 4.57.0
- All required audio/ML packages

### Test Both FA2 and FA3

Once setup completes, verify both versions work:

```bash
# Activate environment
conda activate suno_env_auto

# Test both FA2 and FA3
python ~/projects/ml_build_env/test_fa2_fa3.py
```

**Expected output:**
- ✓ FA2 and FA3 both import successfully
- ✓ Both can perform inference
- ✓ Output shapes match: `torch.Size([2, 128, 8, 64])`

### Run Full Inference Test

For a complete end-to-end test:

```bash
# Run inference test (uses GPU 7 by default)
python ~/projects/ml_build_env/test_fa2_inference.py
```

To use a different GPU, edit line 13 of `test_fa2_inference.py`:
```python
os.environ["CUDA_VISIBLE_DEVICES"] = "7"  # Change to your GPU
```

**Expected output:**
- Model loads successfully
- Generates ~30s song from lyrics
- Saves to `generated_song_fa2.mp3`
- Takes ~20 seconds for generation

### Check All Environments

List all suno environments and their versions:
```bash
source activate_suno.sh
```

## 📦 Prerequisites

- **OS**: Linux (Ubuntu 20.04+ recommended)
- **Python**: 3.10.x (3.10.15 preferred)
- **CUDA**: 12.4 (for GPU support)
- CUDA Toolkit 12.4 (for GPU support)
- Conda (Anaconda/Miniconda)
- `glockenspiel` repository (clone from internal repo to `~/projects/glockenspiel`)
- sox - For audio processing

### Optional Software
- [UV](https://github.com/astral-sh/uv) - For 10-100x faster pip installations

## 🔧 Installation Methods

### Recommended: Flash Attention v2 Setup (Fixed Commit)

Use the setup script from Quick Start section above for the most reliable Flash Attention v2 installation with a tested commit hash.

### Alternative: Legacy Setup Scripts

**Note:** These older scripts may not include the latest fixes. Use the FA2 fixed commit script instead.

<details>
<summary>Click to expand legacy setup options</summary>

#### Conda Setup
```bash
# Basic usage
bash setup_conda_env.sh [environment_name]

# Custom environment name
bash setup_conda_env.sh my_suno_env
```

#### Virtualenv Setup
```bash
# Basic usage
bash setup_virtualenv.sh [environment_name]

# Custom environment name
bash setup_virtualenv.sh my_venv
```

</details>

## ✅ Verification & Testing

### Automated Verification

The setup script automatically runs verification tests at the end. You can also manually verify:

```bash
# Activate environment
conda activate your_env_name

# Run verification script (if available)
python verify_imports.py
```

### Manual Tests
```bash
# Quick verification
python -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
python -c "import flash_attn; print(f'Flash Attention: {flash_attn.__version__}')"
python -c "import suno_utils; print('Suno Utils OK')"
```

### Inference Test

Run the full inference test (recommended):
```bash
python ~/projects/ml_build_env/test_fa2_inference.py
```

This tests:
- Model loading
- Flash Attention integration
- Audio generation pipeline
- Full end-to-end workflow

### Training Test

Run a minimal training test to verify the environment can train models:
```bash
# Activate environment
conda activate your_env_name

# Run training test (will prompt to run actual training)
bash ~/projects/ml_build_env/test_mini_train.sh
```

This tests:
- Environment setup (Python, packages, CUDA)
- Core package imports (PyTorch, Transformers, DeepSpeed, etc.)
- Minimal training run (5 iterations, ~2 minutes)
- Model initialization and forward/backward pass
- Checkpoint saving
- SLURM script compatibility

**Expected results:**
- ✓ All imports successful
- ✓ Training completes 5 iterations
- ✓ Validation loss decreases
- ✓ Uses ~60-80GB GPU memory on H100
- ✓ Logs saved to `test_logs/`

## 🛠 Environment Management

### List All Environments

Check all suno environments and their configurations:

```bash
source activate_suno.sh
```

This shows a table with:
- Environment names
- PyTorch versions
- Flash Attention v2 versions
- Flash Attention v3 availability

### Activate Environment

```bash
# Activate any suno environment
source activate_suno.sh suno_env_auto

# Or use conda directly
conda activate suno_env_auto
```

### Remove Environment

```bash
# Remove conda environment
conda env remove -n suno_env_auto -y
```

## 🔍 Troubleshooting

### Common Issues

#### 1. CUDA Not Available
```bash
# Check NVIDIA driver
nvidia-smi

# Verify CUDA version
nvcc --version

# Reinstall PyTorch with correct CUDA
pip uninstall torch torchvision torchaudio
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
```

#### 2. Flash Attention Build Fails
```bash
# Install from pre-built wheel (if available)
pip install flash-attn --no-build-isolation

# Or build with limited jobs to avoid OOM
export MAX_JOBS=2
cd flash-attention && python setup.py install
```

#### 3. Suno Utils Not Found
```bash
# Verify path exists
ls ~/projects/glockenspiel/suno_utils

# Reinstall
pip install -e ~/projects/glockenspiel/suno_utils
```

#### 4. Import Errors
```bash
# Check package is installed
pip list | grep package_name

# Reinstall specific package
pip install --force-reinstall package_name==version
```

#### 5. Memory Issues During Build
```bash
# Limit parallel compilation jobs
export MAX_JOBS=1

# Increase swap if needed
sudo fallocate -l 16G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile
```

### Environment Variables

Set these before running scripts:

```bash
# GPU selection
export CUDA_VISIBLE_DEVICES=0

# Compilation settings
export MAX_JOBS=4
export TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;9.0"

# Cache directories
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER
export PIP_CACHE_DIR=/mnt/localdisk/.pip_cache_$USER
```

## 📦 Package List

### Core ML Frameworks
- **PyTorch** 2.6.0 + CUDA 12.4
- **Transformers** 4.57.0
- **Flash Attention v3** (H100/Hopper optimized) - from hopper directory
- **Flash Attention v2** 2.8.1 (commit c5b0c631 for compatibility)
- **DeepSpeed** 0.18.0
- **PyTorch Lightning** 2.5.5
- **Weights & Biases** 0.22.2

### Audio Processing
- **nnAudio** 0.3.3
- **Auraloss** (unknown version)
- **Encodec** 0.1.1
- **Sox** (via conda-forge)

### NLP & Text
- **G2P-en** (unknown version)
- **Phonemizer** 3.3.0
- **SentencePiece** 0.2.0
- **TikToken** (unknown version)
- **Better Profanity** 0.7.0

### Utilities
- **Einops** 0.8.1
- **Torch SDE** 0.2.6
- **Modal** 1.0.1
- **Ninja** 1.13.0
- **Suno Utils** 0.0.1a0 (local package)

## 📁 File Structure

```
ml_build_env/
├── README.md                                     # This file
├── activate_suno.sh                              # Unified environment activation script
│
├── Setup Scripts
│   ├── setup_suno_env_fa2_fa3_fixed_commit.sh   # Recommended: FA3 + FA2 (both)
│   ├── setup_suno_env_fa2_fa3_fixed_commit.py   # Python version of above
│   ├── setup_suno_env_fa2_fixed_commit.sh       # FA2 only with fixed commit
│   ├── setup_suno_env_fa3.sh                    # FA3 only (H100 optimized)
│   ├── setup_suno_env_newer.sh                  # FA2 with newer PyTorch
│   └── setup_suno_env_fa2.sh                    # Legacy FA2 setup
│
├── Testing
│   ├── test_fa2_fa3.py                          # Test both FA2 and FA3
│   ├── test_fa2_inference.py                    # Full inference test
│   ├── test_mini_train.sh                       # Minimal training test
│   └── verify_imports.py                        # Package import verification
│
└── Logs
    └── tmp/                                      # Setup logs
```

## 🔄 Updating Environments

To update packages in an existing environment:

```bash
# Activate environment
conda activate your_env_name

# Update specific packages
pip install --upgrade package_name

# Or update PyTorch
pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# Verify after updates
python -c "import torch; import flash_attn; print('OK')"
```

**Note:** Updating Flash Attention requires rebuilding from source. It's often easier to create a new environment.

## 📝 Notes

1. **Flash Attention Dual Installation**: The recommended setup script installs both FA3 (H100 optimized) and FA2 (commit `c5b0c631` for compatibility). Both can coexist and perform inference in the same environment.

2. **GPU Requirements**: Requires NVIDIA GPU with compute capability 7.0+ (V100, A100, H100, etc.). H100/H800 is required for FA3 to work optimally.

3. **Setup Time**: Full environment setup with both FA versions takes 3-4 hours including:
   - Conda environment creation: ~2-5 min
   - Suno utils installation: ~10-15 min
   - PyTorch installation: ~3-5 min
   - Flash Attention v3 build: ~2.5 hours (H100 optimized)
   - Flash Attention v2 build: ~30-60 min
   - Additional packages: ~2-5 min

4. **FA2 vs FA3**: Use FA2 for compatibility with existing code, FA3 for maximum performance on H100. Both use the same API (`flash_attn_func`).

5. **Glockenspiel Dependency**: Make sure glockenspiel repository is on latest main branch before creating environments.

6. **Custom Paths**: Update paths in scripts if your setup differs from `~/projects/` structure.

## 🤝 Support

For issues or questions:
1. Check the troubleshooting section
2. Review setup logs in `setup_*.log`
3. Run verification scripts to identify missing components
4. Check frozen environment snapshots for working configurations

## 📜 License

Internal use only. Contains references to proprietary Suno packages and configurations.