setup.py

"""
setup.py - Run this at the start of app.py to ensure proper Triton installation
Add: import setup  # at the top of app.py after the docstring
"""

import subprocess
import sys
import os

def fix_triton_installation():
    """Fix Triton for MX format by running install.sh if needed."""
    try:
        # Check if we have the right triton
        import triton_kernels
        from triton.tools.ragged_tma import load_ragged, store_ragged
        print("✓ Triton already properly configured for MX format")
        return True
    except ImportError:
        print("Triton not properly configured for MX format")
        print("Running install.sh to fix dependencies...")
        
        # Check if install.sh exists
        if os.path.exists("install.sh"):
            try:
                # Make it executable and run it
                subprocess.check_call(["chmod", "+x", "install.sh"])
                subprocess.check_call(["./install.sh"])
                print("✓ Dependencies installed via install.sh")
                return True
            except subprocess.CalledProcessError as e:
                print(f"Error running install.sh: {e}")
        else:
            print("install.sh not found - trying direct pip fix...")
            
            # Fallback: run key commands directly
            try:
                # Clean and reinstall triton
                subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "triton", "triton_kernels"], 
                             capture_output=True)
                
                # Install nightly triton
                subprocess.check_call([
                    sys.executable, "-m", "pip", "install", "--upgrade",
                    "--index-url", "https://download.pytorch.org/whl/nightly/cu121",
                    "triton"
                ])
                
                # Install triton_kernels
                subprocess.check_call([
                    sys.executable, "-m", "pip", "install",
                    "git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernels"
                ])
                
                print("✓ Triton fixed via direct installation")
                return True
            except Exception as e:
                print(f"Failed to fix Triton: {e}")
                return False

# Auto-run on import
if __name__ != "__main__":
    fix_triton_installation()