web_app/gpu_init.py

"""
GPU initialization module that must be imported BEFORE any SpaCy modules.
This ensures GPU is properly configured before SpaCy loads any models.
"""

import os
import logging

logger = logging.getLogger(__name__)

def initialize_gpu_environment():
    """
    Initialize GPU environment variables and settings before SpaCy import.
    This function should be called at the very beginning of the application.
    """
    # Set environment variables BEFORE any imports
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    os.environ['SPACY_PREFER_GPU'] = '1'
    
    try:
        import torch
        
        if torch.cuda.is_available():
            # Force CUDA initialization
            torch.cuda.init()
            
            # Set default device
            torch.cuda.set_device(0)
            
            # Log GPU info
            gpu_name = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
            
            logger.info(f"GPU initialized: {gpu_name} ({gpu_memory:.1f} GB)")
            logger.info(f"CUDA version: {torch.version.cuda}")
            logger.info(f"PyTorch version: {torch.__version__}")
            
            # Pre-configure SpaCy for GPU
            import spacy
            try:
                # Try require_gpu first for strong enforcement
                spacy.require_gpu(gpu_id=0)
                logger.info("SpaCy GPU enforced with require_gpu()")
            except Exception as e:
                # Fallback to prefer_gpu
                logger.warning(f"require_gpu() failed: {e}, using prefer_gpu()")
                gpu_id = spacy.prefer_gpu(gpu_id=0)
                if gpu_id is not False:
                    logger.info(f"SpaCy GPU enabled with prefer_gpu(): device {gpu_id}")
                else:
                    logger.error("SpaCy GPU initialization failed!")
            
            return True
        else:
            logger.info("No CUDA device available - running on CPU")
            return False
            
    except ImportError:
        logger.info("PyTorch not installed - GPU support unavailable")
        return False
    except Exception as e:
        logger.error(f"GPU initialization error: {e}")
        return False

# Initialize GPU on module import
GPU_AVAILABLE = initialize_gpu_environment()