Spaces:
Building
Building
File size: 2,347 Bytes
bb65e54 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
"""
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()
|