#!/bin/bash
# Activation script for suno_env_fa2_fa3_pt29 environment with Flash Attention v3 + v2 and PyTorch 2.9.0

# Source conda
source "$(conda info --base)/etc/profile.d/conda.sh"

# Activate environment
conda activate suno_env_fa2_fa3_pt29

# Set environment variables
export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
export OMP_NUM_THREADS=1
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER
export PYTHONPATH="/tmp/flash-attention-v3-suno_env_fa2_fa3_pt29/hopper:$PYTHONPATH"

echo "Environment suno_env_fa2_fa3_pt29 activated (Flash Attention v3 + v2, PyTorch 2.9.0)!"
echo "Python: $(which python)"
echo "PyTorch: $(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo 'Not available')"
echo "CUDA available: $(python -c 'import torch; print(torch.cuda.is_available())' 2>/dev/null || echo 'Unknown')"
echo "Flash Attention v3 path: /tmp/flash-attention-v3-suno_env_fa2_fa3_pt29"
echo "Flash Attention v2 path: /tmp/flash-attention-suno_env_fa2_fa3_pt29"
echo "Suno Utils path: /home/victor/glockenspiel/suno_utils"

# Test Flash Attention versions
echo ""
echo "Flash Attention Status:"
python -c "
try:
    import flash_attn_interface
    print('  ✓ FA3 (H100 optimized): Available')
except ImportError:
    print('  ✗ FA3: Not available')

try:
    import flash_attn
    print('  ✓ FA2 (commit 7b0bfcc3d1f69786f0c4277c582ad58acdfb297d): Available')
except ImportError:
    print('  ✗ FA2: Not available')
" 2>/dev/null