setup_environment.sh

# Function definitions
# print_header() {
#     echo "=== $1 ==="
#     echo
# }

# print_section() {
#     echo ">>> $1"
# }

# print_step() {
#     echo ">> $1"
# }

# print_success() {
#     echo "✓ $1"
# }

# print_warning() {
#     echo "⚠ $1"
# }

# print_error() {
#     echo "✗ $1"
# }

# Function to detect GPUs

# Function to create environment file
create_environment_file() {
    print_section "Creating Environment File"

    # Create environment file
    print_step "Creating environment file..."

    # Create .mlstack_env file in home directory
    cat > $HOME/.mlstack_env << EOF
# ML Stack Environment File
# Created by ML Stack Environment Setup Script

# GPU Selection
export HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES
export CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES
export PYTORCH_ROCM_DEVICE=$PYTORCH_ROCM_DEVICE

# Performance Settings
export HSA_ENABLE_SDMA=0
export GPU_MAX_HEAP_SIZE=100
export GPU_MAX_ALLOC_PERCENT=100
export HSA_TOOLS_LIB=1

# MIOpen Settings
export MIOPEN_DEBUG_CONV_IMPLICIT_GEMM=1
export MIOPEN_FIND_MODE=3
export MIOPEN_FIND_ENFORCE=3

# PyTorch Settings
export TORCH_CUDA_ARCH_LIST="7.0;8.0;9.0"
export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:512"

# MPI Settings
export OMPI_MCA_opal_cuda_support=true
export OMPI_MCA_pml_ucx_opal_cuda_support=true
export OMPI_MCA_btl_openib_allow_ib=true
export OMPI_MCA_btl_openib_warn_no_device_params_found=0
export OMPI_MCA_coll_hcoll_enable=0
export OMPI_MCA_pml=ucx
export OMPI_MCA_osc=ucx
export OMPI_MCA_btl=^openib,uct

# Path Settings
export PATH=\$PATH:/opt/rocm/bin:/opt/rocm/hip/bin
export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/opt/rocm/lib:/opt/rocm/hip/lib
EOF

    # Add source to .bashrc if not already there
    if ! grep -q "source \$HOME/.mlstack_env" $HOME/.bashrc; then
        echo -e "\n# Source ML Stack environment" >> $HOME/.bashrc
        echo "source \$HOME/.mlstack_env" >> $HOME/.bashrc
    fi

    # Source the file
    source $HOME/.mlstack_env

    print_success "Environment file created successfully"
    print_step "Environment file: $HOME/.mlstack_env"
    print_step "The environment file has been added to your .bashrc file."
    print_step "To apply the changes, run: source $HOME/.bashrc"

    return 0
}

# Main function
main() {
    print_header "ML Stack Environment Setup"

    # Detect GPUs
    detect_gpus
    if [ $? -ne 0 ]; then
        print_error "GPU detection failed. Exiting."
        exit 1
    fi

    # Create environment file
    create_environment_file
    if [ $? -ne 0 ]; then
        print_error "Environment file creation failed. Exiting."
        exit 1
    fi


    print_header "ML Stack Environment Setup Complete"
    print_step "To apply the changes, run: source $HOME/.bashrc"

    return 0
}

# Run main function
main