Spaces:
Sleeping
Sleeping
File size: 3,619 Bytes
e06a21d 1dd0906 e06a21d | 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 | #!/usr/bin/env python3
"""
Startup script to handle environment setup and pre-checks.
This script ensures proper directory permissions and environment setup before starting the app.
"""
import os
import sys
import tempfile
from pathlib import Path
def setup_environment():
"""Setup environment variables and directories."""
# Set cache directories to writable locations
cache_dirs = {
'SENTENCE_TRANSFORMERS_HOME': ['/tmp/sentence_transformers', '/data/sentence_transformers'],
'VECTOR_PERSIST_DIR': ['/tmp/vector_db', '/data/vector_db'],
'HF_HOME': ['/tmp/huggingface', '/data/huggingface'],
'TRANSFORMERS_CACHE': ['/tmp/transformers', '/data/transformers'],
}
for env_var, candidates in cache_dirs.items():
if env_var not in os.environ:
for candidate in candidates:
try:
Path(candidate).mkdir(parents=True, exist_ok=True)
# Test write permission
test_file = Path(candidate) / "test_write"
test_file.write_text("test")
test_file.unlink()
os.environ[env_var] = candidate
print(f"✅ Set {env_var} to {candidate}")
break
except (OSError, PermissionError):
continue
else:
# Fallback to system temp
temp_dir = os.path.join(tempfile.gettempdir(), env_var.lower())
try:
Path(temp_dir).mkdir(parents=True, exist_ok=True)
os.environ[env_var] = temp_dir
print(f"⚠️ Fallback: Set {env_var} to {temp_dir}")
except Exception as e:
print(f"❌ Could not set {env_var}: {e}")
# Ensure tokenizer libs don't attempt parallelism post-fork (prevents hangs)
os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false')
# Cap thread usage for BLAS backends to avoid CPU thrashing on shared runners
os.environ.setdefault('OMP_NUM_THREADS', '1')
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
os.environ.setdefault('MKL_NUM_THREADS', '1')
os.environ.setdefault('NUMEXPR_NUM_THREADS', '1')
def check_permissions():
"""Check and report on directory permissions."""
dirs_to_check = [
os.environ.get('SENTENCE_TRANSFORMERS_HOME', '/tmp/sentence_transformers'),
os.environ.get('VECTOR_PERSIST_DIR', '/tmp/vector_db'),
'/tmp',
'.',
]
print("\n📋 Directory Permissions Check:")
for dir_path in dirs_to_check:
if not dir_path:
continue
try:
path = Path(dir_path)
if not path.exists():
path.mkdir(parents=True, exist_ok=True)
# Test write permission
test_file = path / "test_permissions"
test_file.write_text("test")
test_file.unlink()
print(f"✅ {dir_path} - writable")
except Exception as e:
print(f"❌ {dir_path} - not writable: {e}")
def main():
"""Main startup routine."""
print("🚀 Starting CVE Fact Checker setup...")
setup_environment()
check_permissions()
print("\n🎯 Environment setup complete!")
print("📍 Key directories:")
for key in ['SENTENCE_TRANSFORMERS_HOME', 'VECTOR_PERSIST_DIR']:
print(f" {key}: {os.environ.get(key, 'Not set')}")
return True
if __name__ == "__main__":
success = main()
sys.exit(0 if success else 1)
|