File size: 7,096 Bytes
3a660a3 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 |
"""
Environment Detection Utility
Detects GPU availability, HuggingFace Space environment, and system capabilities
"""
import os
import platform
import logging
from typing import Dict, Any, Optional
logger = logging.getLogger(__name__)
class EnvironmentDetector:
"""Detect runtime environment and capabilities"""
def __init__(self):
self._gpu_available: Optional[bool] = None
self._is_huggingface: Optional[bool] = None
self._transformers_available: Optional[bool] = None
self._torch_available: Optional[bool] = None
def is_huggingface_space(self) -> bool:
"""Detect if running on HuggingFace Space"""
if self._is_huggingface is None:
# Check for HF Space environment variables
self._is_huggingface = bool(
os.getenv("SPACE_ID") or
os.getenv("SPACE_AUTHOR_NAME") or
os.getenv("SPACE_HOST")
)
return self._is_huggingface
def has_gpu(self) -> bool:
"""Detect if GPU is available"""
if self._gpu_available is None:
self._gpu_available = False
try:
import torch
self._gpu_available = torch.cuda.is_available()
if self._gpu_available:
gpu_name = torch.cuda.get_device_name(0)
logger.info(f"✅ GPU detected: {gpu_name}")
else:
logger.info("ℹ️ No GPU detected - using CPU")
except ImportError:
logger.info("ℹ️ PyTorch not installed - assuming no GPU")
self._gpu_available = False
except Exception as e:
logger.warning(f"Error detecting GPU: {e}")
self._gpu_available = False
return self._gpu_available
def is_torch_available(self) -> bool:
"""Check if PyTorch is installed"""
if self._torch_available is None:
try:
import torch
self._torch_available = True
logger.info(f"✅ PyTorch {torch.__version__} available")
except ImportError:
self._torch_available = False
logger.info("ℹ️ PyTorch not installed")
return self._torch_available
def is_transformers_available(self) -> bool:
"""Check if Transformers library is installed"""
if self._transformers_available is None:
try:
import transformers
self._transformers_available = True
logger.info(f"✅ Transformers {transformers.__version__} available")
except ImportError:
self._transformers_available = False
logger.info("ℹ️ Transformers not installed")
return self._transformers_available
def should_use_ai_models(self) -> bool:
"""
Determine if AI models should be used
Only use if:
- Running on HuggingFace Space, OR
- Transformers is installed AND (GPU available OR explicitly enabled)
"""
if self.is_huggingface_space():
logger.info("✅ HuggingFace Space detected - AI models will be used")
return True
if not self.is_transformers_available():
logger.info("ℹ️ Transformers not available - using fallback mode")
return False
# If transformers installed but not HF Space, check GPU or explicit flag
use_ai = os.getenv("USE_AI_MODELS", "").lower() == "true" or self.has_gpu()
if use_ai:
logger.info("✅ AI models enabled (GPU or USE_AI_MODELS=true)")
else:
logger.info("ℹ️ AI models disabled (no GPU, set USE_AI_MODELS=true to force)")
return use_ai
def get_device(self) -> str:
"""Get the device to use for AI models"""
if self.has_gpu():
return "cuda"
return "cpu"
def get_environment_info(self) -> Dict[str, Any]:
"""Get comprehensive environment information"""
info = {
"platform": platform.system(),
"python_version": platform.python_version(),
"is_huggingface_space": self.is_huggingface_space(),
"torch_available": self.is_torch_available(),
"transformers_available": self.is_transformers_available(),
"gpu_available": self.has_gpu(),
"device": self.get_device() if self.is_torch_available() else "none",
"should_use_ai": self.should_use_ai_models()
}
# Add GPU details if available
if self.has_gpu():
try:
import torch
info["gpu_name"] = torch.cuda.get_device_name(0)
info["gpu_count"] = torch.cuda.device_count()
info["cuda_version"] = torch.version.cuda
except:
pass
# Add HF Space info if available
if self.is_huggingface_space():
info["space_id"] = os.getenv("SPACE_ID", "unknown")
info["space_author"] = os.getenv("SPACE_AUTHOR_NAME", "unknown")
return info
def log_environment(self):
"""Log environment information"""
info = self.get_environment_info()
logger.info("=" * 70)
logger.info("🔍 ENVIRONMENT DETECTION:")
logger.info(f" Platform: {info['platform']}")
logger.info(f" Python: {info['python_version']}")
logger.info(f" HuggingFace Space: {'Yes' if info['is_huggingface_space'] else 'No'}")
logger.info(f" PyTorch: {'Yes' if info['torch_available'] else 'No'}")
logger.info(f" Transformers: {'Yes' if info['transformers_available'] else 'No'}")
logger.info(f" GPU: {'Yes' if info['gpu_available'] else 'No'}")
if info['gpu_available'] and 'gpu_name' in info:
logger.info(f" GPU Name: {info['gpu_name']}")
logger.info(f" Device: {info['device']}")
logger.info(f" AI Models: {'Enabled' if info['should_use_ai'] else 'Disabled (using fallback)'}")
logger.info("=" * 70)
# Global instance
_env_detector = EnvironmentDetector()
def get_environment_detector() -> EnvironmentDetector:
"""Get global environment detector instance"""
return _env_detector
def is_huggingface_space() -> bool:
"""Quick check if running on HuggingFace Space"""
return _env_detector.is_huggingface_space()
def has_gpu() -> bool:
"""Quick check if GPU is available"""
return _env_detector.has_gpu()
def should_use_ai_models() -> bool:
"""Quick check if AI models should be used"""
return _env_detector.should_use_ai_models()
def get_device() -> str:
"""Get device for AI models"""
return _env_detector.get_device()
__all__ = [
'EnvironmentDetector',
'get_environment_detector',
'is_huggingface_space',
'has_gpu',
'should_use_ai_models',
'get_device'
]
|