371 lines
12 KiB
Python
371 lines
12 KiB
Python
|
"""VRAM + Resonance Manager.
|
||
|
|
|
||
|
|
Combines computational VRAM limits with symbolic resonance:
|
||
|
|
- Monitors GPU VRAM (8GB GTX1080)
|
||
|
|
- Adjusts model loading based on glyph resonance
|
||
|
|
- Prevents crashes from simultaneous Forge + Janus
|
||
|
|
- Dynamic VRAM budgeting from glyph activation
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
from dual_layer.vram_manager import VRAMManager
|
||
|
|
|
||
|
|
manager = VRAMManager()
|
||
|
|
if manager.can_activate_glyph(glyph_routing_result):
|
||
|
|
manager.activate(glyph_routing_result)
|
||
|
|
"""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import Dict, List, Any, Optional, Tuple
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from datetime import datetime
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
# VRAM constants (GTX 1080: 8GB)
|
||
|
|
MAX_VRAM = 8.0
|
||
|
|
WARNING_THRESHOLD = 6.5
|
||
|
|
CRITICAL_THRESHOLD = 7.5
|
||
|
|
VRAM_WARNING_GB = 6.5
|
||
|
|
VRAM_CRITICAL_GB = 7.5
|
||
|
|
VRAM_TOTAL_GB = 8.0
|
||
|
|
|
||
|
|
# Model VRAM estimates
|
||
|
|
MODEL_VRAM_ESTIMATES: Dict[str, float] = {
|
||
|
|
"llama": 2.0, # Llama 7B ~2GB
|
||
|
|
"forge": 4.5, # Stable Diffusion XL ~4.5GB
|
||
|
|
"janus": 5.0, # Janus-Pro-7B ~5GB
|
||
|
|
"google_ai": 1.5, # Google AI API (minimal local)
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class GlyphActivation:
|
||
|
|
"""Active glyph reservation."""
|
||
|
|
glyph_id: str
|
||
|
|
specialized_type: str
|
||
|
|
model: str
|
||
|
|
vram_budget: float
|
||
|
|
resonance_score: float
|
||
|
|
power_boost: float
|
||
|
|
activated_at: datetime
|
||
|
|
priority: float
|
||
|
|
|
||
|
|
|
||
|
|
class VRAMManager:
|
||
|
|
"""Manages VRAM + resonance for dual-layer system."""
|
||
|
|
|
||
|
|
def __init__(self, total_vram: float = VRAM_TOTAL_GB):
|
||
|
|
self.total_vram = total_vram
|
||
|
|
self.active_glyphs: Dict[str, GlyphActivation] = {}
|
||
|
|
self.vram_usage: float = 0.0
|
||
|
|
self._lock = asyncio.Lock() # Async lock for concurrent safety
|
||
|
|
|
||
|
|
# Model state tracking
|
||
|
|
self.loaded_models: Dict[str, bool] = {
|
||
|
|
"llama": False,
|
||
|
|
"forge": False,
|
||
|
|
"janus": False,
|
||
|
|
"google_ai": False,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Critical rule: NEVER run Forge + Janus simultaneously
|
||
|
|
self._forge_active = False
|
||
|
|
self._janus_active = False
|
||
|
|
|
||
|
|
async def get_vram_status(self) -> Dict[str, Any]:
|
||
|
|
"""Get current VRAM status."""
|
||
|
|
async with self._lock:
|
||
|
|
return {
|
||
|
|
"total_vram_gb": self.total_vram,
|
||
|
|
"used_vram_gb": self.vram_usage,
|
||
|
|
"available_vram_gb": self.total_vram - self.vram_usage,
|
||
|
|
"usage_percent": (self.vram_usage / self.total_vram) * 100,
|
||
|
|
"active_glyphs": len(self.active_glyphs),
|
||
|
|
"warning": self.vram_usage >= VRAM_WARNING_GB,
|
||
|
|
"critical": self.vram_usage >= VRAM_CRITICAL_GB,
|
||
|
|
"loaded_models": self.loaded_models,
|
||
|
|
"forge_active": self._forge_active,
|
||
|
|
"janus_active": self._janus_active,
|
||
|
|
}
|
||
|
|
|
||
|
|
def can_activate_glyph(
|
||
|
|
self,
|
||
|
|
glyph_id: str,
|
||
|
|
model: str,
|
||
|
|
vram_budget: float,
|
||
|
|
priority: float
|
||
|
|
) -> Tuple[bool, str]:
|
||
|
|
"""Check if glyph can be activated without VRAM crash.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
glyph_id: Glyph identifier
|
||
|
|
model: Model to use (llama, forge, janus, google_ai)
|
||
|
|
vram_budget: Requested VRAM budget
|
||
|
|
priority: Glyph priority (higher = more authority)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
(can_activate, reason)
|
||
|
|
"""
|
||
|
|
# Check critical VRAM
|
||
|
|
if self.vram_usage >= VRAM_CRITICAL_GB:
|
||
|
|
return False, f"Critical VRAM: {self.vram_usage:.2f}GB used"
|
||
|
|
|
||
|
|
# Check Forge + Janus mutex
|
||
|
|
if model == "forge" and self._janus_active:
|
||
|
|
return False, "Forge cannot run while Janus is active (VRAM crash risk)"
|
||
|
|
|
||
|
|
if model == "janus" and self._forge_active:
|
||
|
|
return False, "Janus cannot run while Forge is active (VRAM crash risk)"
|
||
|
|
|
||
|
|
# Check available VRAM
|
||
|
|
projected_usage = self.vram_usage + vram_budget
|
||
|
|
if projected_usage > self.total_vram:
|
||
|
|
# Check if we can deactivate lower-priority glyphs
|
||
|
|
can_free = self._can_free_vram_for(
|
||
|
|
vram_budget,
|
||
|
|
priority,
|
||
|
|
model
|
||
|
|
)
|
||
|
|
if not can_free:
|
||
|
|
return False, f"Insufficient VRAM: need {vram_budget:.2f}GB, have {self.total_vram - self.vram_usage:.2f}GB available"
|
||
|
|
|
||
|
|
# Check warning threshold
|
||
|
|
if projected_usage >= VRAM_WARNING_GB:
|
||
|
|
logger.warning(
|
||
|
|
f"Glyph {glyph_id} activation will trigger VRAM warning "
|
||
|
|
f"({projected_usage:.2f}GB >= {VRAM_WARNING_GB}GB)"
|
||
|
|
)
|
||
|
|
|
||
|
|
return True, "OK"
|
||
|
|
|
||
|
|
def _can_free_vram_for(
|
||
|
|
self,
|
||
|
|
needed_vram: float,
|
||
|
|
priority: float,
|
||
|
|
model: str
|
||
|
|
) -> bool:
|
||
|
|
"""Check if we can free VRAM by deactivating lower-priority glyphs."""
|
||
|
|
available = self.total_vram - self.vram_usage
|
||
|
|
|
||
|
|
# Find lower-priority glyphs
|
||
|
|
lower_priority_glyphs = [
|
||
|
|
(gid, activation)
|
||
|
|
for gid, activation in self.active_glyphs.items()
|
||
|
|
if activation.priority < priority
|
||
|
|
]
|
||
|
|
|
||
|
|
# Sort by priority (lowest first)
|
||
|
|
lower_priority_glyphs.sort(key=lambda x: x[1].priority)
|
||
|
|
|
||
|
|
# Calculate if deactivating would free enough
|
||
|
|
potential_free = available
|
||
|
|
for _, activation in lower_priority_glyphs:
|
||
|
|
potential_free += activation.vram_budget
|
||
|
|
if potential_free >= needed_vram:
|
||
|
|
return True
|
||
|
|
|
||
|
|
return False
|
||
|
|
|
||
|
|
async def activate_glyph(
|
||
|
|
self,
|
||
|
|
glyph_id: str,
|
||
|
|
specialized_type: str,
|
||
|
|
model: str,
|
||
|
|
vram_budget: float,
|
||
|
|
resonance_score: float,
|
||
|
|
power_boost: float,
|
||
|
|
priority: float
|
||
|
|
) -> bool:
|
||
|
|
"""Activate a glyph (reserve VRAM).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
glyph_id: Glyph identifier
|
||
|
|
specialized_type: Glyph specialized type
|
||
|
|
model: Model to use
|
||
|
|
vram_budget: VRAM budget
|
||
|
|
resonance_score: Resonance score (0-100)
|
||
|
|
power_boost: Power boost multiplier
|
||
|
|
priority: Priority level
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
True if activated, False if failed
|
||
|
|
"""
|
||
|
|
async with self._lock:
|
||
|
|
# Check again under lock
|
||
|
|
can_activate, reason = self.can_activate_glyph(
|
||
|
|
glyph_id, model, vram_budget, priority
|
||
|
|
)
|
||
|
|
|
||
|
|
if not can_activate:
|
||
|
|
logger.error(f"Cannot activate {glyph_id}: {reason}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
# Deactivate lower-priority glyphs if needed
|
||
|
|
self._deactivate_lower_priority(priority, vram_budget)
|
||
|
|
|
||
|
|
# Create activation record
|
||
|
|
activation = GlyphActivation(
|
||
|
|
glyph_id=glyph_id,
|
||
|
|
specialized_type=specialized_type,
|
||
|
|
model=model,
|
||
|
|
vram_budget=vram_budget,
|
||
|
|
resonance_score=resonance_score,
|
||
|
|
power_boost=power_boost,
|
||
|
|
activated_at=datetime.now(),
|
||
|
|
priority=priority
|
||
|
|
)
|
||
|
|
|
||
|
|
# Track model loading
|
||
|
|
if not self.loaded_models.get(model, False):
|
||
|
|
logger.info(f"Loading model: {model} (estimated {MODEL_VRAM_ESTIMATES.get(model, 0):.1f}GB)")
|
||
|
|
self.loaded_models[model] = True
|
||
|
|
|
||
|
|
# Track Forge/Janus mutex
|
||
|
|
if model == "forge":
|
||
|
|
self._forge_active = True
|
||
|
|
elif model == "janus":
|
||
|
|
self._janus_active = True
|
||
|
|
|
||
|
|
# Reserve VRAM
|
||
|
|
self.active_glyphs[glyph_id] = activation
|
||
|
|
self.vram_usage += vram_budget
|
||
|
|
|
||
|
|
logger.info(
|
||
|
|
f"✅ Activated glyph {glyph_id} ({specialized_type}) "
|
||
|
|
f"→ {model} model, {vram_budget:.2f}GB VRAM, "
|
||
|
|
f"resonance={resonance_score:.1f}, boost={power_boost:.2f}x"
|
||
|
|
)
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
async def deactivate_glyph(self, glyph_id: str) -> bool:
|
||
|
|
"""Deactivate a glyph (release VRAM).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
glyph_id: Glyph identifier
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
True if deactivated, False if not found
|
||
|
|
"""
|
||
|
|
async with self._lock:
|
||
|
|
if glyph_id not in self.active_glyphs:
|
||
|
|
return False
|
||
|
|
|
||
|
|
activation = self.active_glyphs.pop(glyph_id)
|
||
|
|
self.vram_usage -= activation.vram_budget
|
||
|
|
|
||
|
|
# Track model unloading
|
||
|
|
model = activation.model
|
||
|
|
if self.loaded_models.get(model, False):
|
||
|
|
# Check if any other glyphs use this model
|
||
|
|
model_users = sum(
|
||
|
|
1 for a in self.active_glyphs.values()
|
||
|
|
if a.model == model
|
||
|
|
)
|
||
|
|
if model_users == 0:
|
||
|
|
logger.info(f"Unloading model: {model}")
|
||
|
|
self.loaded_models[model] = False
|
||
|
|
|
||
|
|
# Track Forge/Janus mutex
|
||
|
|
if model == "forge":
|
||
|
|
self._forge_active = False
|
||
|
|
elif model == "janus":
|
||
|
|
self._janus_active = False
|
||
|
|
|
||
|
|
logger.info(
|
||
|
|
f"❌ Deactivated glyph {glyph_id} "
|
||
|
|
f"(released {activation.vram_budget:.2f}GB VRAM)"
|
||
|
|
)
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
def _deactivate_lower_priority(
|
||
|
|
self,
|
||
|
|
priority: float,
|
||
|
|
needed_vram: float
|
||
|
|
):
|
||
|
|
"""Deactivate lower-priority glyphs to free VRAM."""
|
||
|
|
available = self.total_vram - self.vram_usage
|
||
|
|
|
||
|
|
if available >= needed_vram:
|
||
|
|
return # No need to deactivate
|
||
|
|
|
||
|
|
# Find and sort lower-priority glyphs
|
||
|
|
lower_priority_glyphs = [
|
||
|
|
(gid, activation)
|
||
|
|
for gid, activation in self.active_glyphs.items()
|
||
|
|
if activation.priority < priority
|
||
|
|
]
|
||
|
|
lower_priority_glyphs.sort(key=lambda x: x[1].priority)
|
||
|
|
|
||
|
|
# Deactivate until enough VRAM is freed
|
||
|
|
for glyph_id, activation in lower_priority_glyphs:
|
||
|
|
self.deactivate_glyph(glyph_id)
|
||
|
|
available += activation.vram_budget
|
||
|
|
|
||
|
|
if available >= needed_vram:
|
||
|
|
logger.info(
|
||
|
|
f"Deactivated {len(lower_priority_glyphs)} lower-priority "
|
||
|
|
f"glyphs to free {needed_vram - (self.total_vram - available):.2f}GB"
|
||
|
|
)
|
||
|
|
break
|
||
|
|
|
||
|
|
def get_active_glyphs(self) -> List[Dict[str, Any]]:
|
||
|
|
"""Get list of active glyphs."""
|
||
|
|
return [
|
||
|
|
{
|
||
|
|
"glyph_id": a.glyph_id,
|
||
|
|
"specialized_type": a.specialized_type,
|
||
|
|
"model": a.model,
|
||
|
|
"vram_budget": a.vram_budget,
|
||
|
|
"resonance_score": a.resonance_score,
|
||
|
|
"power_boost": a.power_boost,
|
||
|
|
"priority": a.priority,
|
||
|
|
"activated_at": a.activated_at.isoformat(),
|
||
|
|
}
|
||
|
|
for a in self.active_glyphs.values()
|
||
|
|
]
|
||
|
|
|
||
|
|
def get_resonance_summary(self) -> Dict[str, Any]:
|
||
|
|
"""Get resonance-based VRAM summary."""
|
||
|
|
if not self.active_glyphs:
|
||
|
|
return {
|
||
|
|
"total_resonance": 0,
|
||
|
|
"average_resonance": 0,
|
||
|
|
"highest_priority_glyph": None,
|
||
|
|
"model_distribution": {},
|
||
|
|
}
|
||
|
|
|
||
|
|
# Calculate resonance metrics
|
||
|
|
total_resonance = sum(a.resonance_score for a in self.active_glyphs.values())
|
||
|
|
avg_resonance = total_resonance / len(self.active_glyphs)
|
||
|
|
|
||
|
|
# Find highest priority
|
||
|
|
highest = max(self.active_glyphs.values(), key=lambda a: a.priority)
|
||
|
|
|
||
|
|
# Model distribution
|
||
|
|
model_counts = {}
|
||
|
|
for a in self.active_glyphs.values():
|
||
|
|
model_counts[a.model] = model_counts.get(a.model, 0) + 1
|
||
|
|
|
||
|
|
return {
|
||
|
|
"total_resonance": total_resonance,
|
||
|
|
"average_resonance": avg_resonance,
|
||
|
|
"highest_priority_glyph": highest.glyph_id,
|
||
|
|
"highest_priority_type": highest.specialized_type,
|
||
|
|
"model_distribution": model_counts,
|
||
|
|
"vram_efficiency": total_resonance / self.vram_usage if self.vram_usage > 0 else 0,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
# Global singleton instance
|
||
|
|
_vram_manager: Optional[VRAMManager] = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_vram_manager() -> VRAMManager:
|
||
|
|
"""Get global VRAM manager instance."""
|
||
|
|
global _vram_manager
|
||
|
|
if _vram_manager is None:
|
||
|
|
_vram_manager = VRAMManager()
|
||
|
|
return _vram_manager
|