""" model_manager.py — Lazy Model Loading, Re-Patching, and PX Metrics =================================================================== Manages model lifecycle on GPU. Lazy loads on first request. Re-patches instead of reloading when switching Peak ↔ Subjective. """ import torch import sys import os import time import importlib import asyncio from typing import Dict, Optional from config import MODEL_REGISTRY from transformers import AutoTokenizer, AutoModelForCausalLM # Preset-Migration (post 2026-06-11). States: BASELINE, ACTIVE_MANIFOLD, # ACTIVE_MANIFOLD_LEAN (kausaler Kern — Crutches weg), ACTIVE_MANIFOLD_RELAY # (LEAN + verstärkbar Selbst-Injektions-Relay, psychomotrik seite15). # Old presets → ACTIVE_MANIFOLD. _VALID_PRESETS = {"BASELINE", "ACTIVE_MANIFOLD", "ACTIVE_MANIFOLD_LEAN", "ACTIVE_MANIFOLD_RELAY"} def _migrate_preset(preset: str) -> str: """Map any old preset name to one of the three valid states.""" if preset in _VALID_PRESETS: return preset return "ACTIVE_MANIFOLD" # gnadenlose Migration class ModelManager: def __init__(self, max_loaded_models: int = 1): self._models: Dict[str, dict] = {} # model_id -> loaded entry self._loading: Dict[str, bool] = {} # model_id -> loading flag self._last_used: Dict[str, float] = {} # model_id -> timestamp self._px_metrics: Dict[str, dict] = {} # model_id -> latest metrics self._busy: set = set() # model_id -> set of active usage self.max_loaded_models = max_loaded_models self._lock = asyncio.Lock() # Global lock for model state changes def lock_model(self, model_id: str): # ... (rest of methods) """Mark a model as busy (in use by a thread).""" self._busy.add(model_id) def unlock_model(self, model_id: str): """Mark a model as free.""" if model_id in self._busy: self._busy.remove(model_id) def is_busy(self, model_id: str) -> bool: """Check if a model is currently in use.""" return model_id in self._busy async def get_model(self, model_id: str, px_subjective: bool = False, px_gamma: float = None, px_routing_mode: str = None, px_config_preset: str = "ACTIVE_MANIFOLD", px_relay_sign: int = None, px_relay_alpha: float = None, px_relay_layer: int = None, quantization: str = None) -> dict: """Get a loaded model, loading lazily if needed. quantization: None → read registry default (gemma3-4b-it ships as int8 because bf16 doesn't fit 12 GB at long prefill; 1b/270m default to "none"). Pass "none" or "int8" explicitly to override. """ async with self._lock: if model_id not in MODEL_REGISTRY: raise ValueError(f"Unknown model: {model_id}") # Resolve quantization=None → registry default if quantization is None: quantization = MODEL_REGISTRY[model_id].get("quantization", "none") if quantization is None: quantization = "none" if model_id in self._models: entry = self._models[model_id] current_subjective = entry.get("px_subjective", False) current_preset = entry.get("px_config_preset") needs_repatch = current_subjective != px_subjective if px_config_preset is not None and current_preset != px_config_preset: needs_repatch = True if px_gamma is not None and entry.get("px_gamma") != px_gamma: needs_repatch = True if px_routing_mode is not None and entry.get("px_routing_mode") != px_routing_mode: needs_repatch = True if px_relay_sign is not None and entry.get("px_relay_sign") != px_relay_sign: needs_repatch = True if px_relay_alpha is not None and entry.get("px_relay_alpha") != px_relay_alpha: needs_repatch = True if px_relay_layer is not None and entry.get("px_relay_layer") != px_relay_layer: needs_repatch = True if needs_repatch: # Safety: wait until model is free before re-patching # A generation might take a while, especially with recursion. print(f"[ModelManager] Waiting for {model_id} to be free for re-patching...") # We are holding self._lock, but we might need to wait for is_busy. # This is tricky because if we wait here, we block ALL other get_model calls. # But re-patching is rare and must be exclusive. while self.is_busy(model_id): # Release control but stay in lock (not ideal for all models, # but necessary for this model) await asyncio.sleep(0.1) self._reapply_patch(model_id, px_subjective, px_gamma, px_routing_mode, px_config_preset, px_relay_sign, px_relay_alpha, px_relay_layer) self._last_used[model_id] = time.time() return entry # Handle auto-unloading before loading new model while len(self._models) >= self.max_loaded_models: # Unload LRU model lru_model = min(self._last_used, key=self._last_used.get) # Safety: wait if model is busy print(f"[ModelManager] Waiting for {lru_model} to be free for unloading...") while self.is_busy(lru_model): await asyncio.sleep(0.1) print(f"[ModelManager] Unloading model {lru_model} to free memory...") self.unload(lru_model) # Lazy load if self._loading.get(model_id): raise RuntimeError(f"Model {model_id} is currently loading") self._loading[model_id] = True try: # Run blocking load in a thread to keep event loop alive loop = asyncio.get_running_loop() entry = await loop.run_in_executor(None, lambda: self._load_model(model_id, px_subjective, px_gamma, px_routing_mode, px_config_preset, px_relay_sign, px_relay_alpha, px_relay_layer, quantization)) self._models[model_id] = entry self._last_used[model_id] = time.time() return entry finally: self._loading[model_id] = False def _load_model(self, model_id: str, px_subjective: bool, px_gamma: float = None, px_routing_mode: str = None, px_config_preset: str = "ACTIVE_MANIFOLD", px_relay_sign: int = None, px_relay_alpha: float = None, px_relay_layer: int = None, quantization: str = None) -> dict: """Load model weights + tokenizer + apply PX patch. quantization: None → read registry default. Registry default for gemma3-4b-it is "int8" (bf16 OOMs on 12 GB at long prefill). """ registry = MODEL_REGISTRY[model_id] if quantization is None: quantization = registry.get("quantization") or "none" hf_id = registry["hf_id"] tok_id = registry["tokenizer_id"] model_type = registry.get("model_type", "gemma3") print(f"[ModelManager] Loading {model_id} from {hf_id} (type={model_type}, subjective={px_subjective}, preset={px_config_preset})...") # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(tok_id) # Apply manual chat template if needed (Gemma3 base model) if registry.get("chat_template_manual"): tokenizer.chat_template = registry["chat_template_manual"] # Load processor for multimodal models (gemma3_conditional / gemma4_conditional). # Text-only models (gemma3, llama) leave processor=None; generators.py rejects # image-bearing requests when processor is None with HTTP 400. processor = None if model_type in ("gemma3_conditional", "gemma4_conditional"): from transformers import AutoProcessor processor = AutoProcessor.from_pretrained(tok_id) print(f"[ModelManager] {model_id} processor loaded for multimodal inputs.") # Load model dtype = getattr(torch, registry["dtype"]) if model_type == "gemma3_conditional": from transformers import Gemma3ForConditionalGeneration model = Gemma3ForConditionalGeneration.from_pretrained( hf_id, torch_dtype=dtype, device_map="auto", ) elif model_type == "gemma4_conditional": from transformers import AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained( hf_id, torch_dtype=dtype, device_map="auto", trust_remote_code=True, ) else: model = AutoModelForCausalLM.from_pretrained( hf_id, torch_dtype=dtype, device_map="auto", ) # Apply weight quantization (Plan 1, Phase D). Default "none" preserves # the legacy bf16/fp16 path. "int8" monkey-patches every nn.Linear with # QuantizedLinear from scratches/4b-image/ — reduces VRAM by ~50% on # 4b (8 GB → ~4.5 GB weights), sufficient to fit long prefills on # 12 GB GPUs that previously OOM'd at the MLP layer. # Quantization happens BEFORE the PX patch so that _px_forward sees # the (already quantized) Linear modules. The patch walks the module # tree normally; our QuantizedLinear subclasses nn.Module transparently. if quantization not in ("none", "int8"): raise ValueError(f"unsupported quantization {quantization!r}; " "supported: 'none', 'int8'") if quantization == "int8": try: from scratches._4b_image.quantize_pipeline import quantize_all_linears except ImportError: # Fallback: sys.path-based import (when called from a different # working dir). scratches/ is two levels deep from this file. import sys, os _scratches_path = os.path.join( os.path.dirname(os.path.abspath(__file__)), "scratches", "4b-image") if _scratches_path not in sys.path: sys.path.insert(0, _scratches_path) from quantize_pipeline import quantize_all_linears n_replaced = quantize_all_linears(model) print(f"[ModelManager] {model_id} quantized int8: {n_replaced} Linears replaced") else: print(f"[ModelManager] {model_id} quantization=none (bf16 path)") # Apply PX patch (skip for unpatched baseline models or if BASELINE preset selected) if registry.get("patch_dir") is not None and px_config_preset != "BASELINE": patch_kwargs = dict(registry["patch_kwargs"]) if px_subjective and registry.get("subjective_kwargs"): patch_kwargs.update(registry["subjective_kwargs"]) # Preset override (migrate old names to ACTIVE_MANIFOLD) if px_config_preset is not None: patch_kwargs["config_preset"] = _migrate_preset(px_config_preset) if px_gamma is not None: patch_kwargs["gamma"] = px_gamma if px_routing_mode is not None: patch_kwargs["routing_mode"] = px_routing_mode # SR-61b: Explicitly pass subjective flag patch_kwargs["subjective_enabled"] = px_subjective # verstärkbar Relay-Parameter (psychomotrik seite15) → patch_kwargs if px_relay_sign is not None: patch_kwargs["relay_sign"] = px_relay_sign if px_relay_alpha is not None: patch_kwargs["relay_alpha"] = px_relay_alpha if px_relay_layer is not None: patch_kwargs["relay_layer"] = px_relay_layer apply_fn = self._get_patch_function(model_id, "apply_px_patch") apply_fn(model, **patch_kwargs) tm = self._resolve_text_model(model) model.tokenizer = tm.tokenizer = tokenizer print(f"[ModelManager] {model_id} loaded and patched successfully.") else: print(f"[ModelManager] {model_id} loaded WITHOUT patch (baseline).") return { "model": model, "tokenizer": tokenizer, "processor": processor, "registry": registry, "px_subjective": px_subjective, "px_gamma": px_gamma, "px_routing_mode": px_routing_mode, "px_config_preset": px_config_preset, "px_relay_sign": px_relay_sign, "px_relay_alpha": px_relay_alpha, "px_relay_layer": px_relay_layer, "model_type": model_type, } def _reapply_patch(self, model_id: str, px_subjective: bool, px_gamma: float = None, px_routing_mode: str = None, px_config_preset: str = "ACTIVE_MANIFOLD", px_relay_sign: int = None, px_relay_alpha: float = None, px_relay_layer: int = None): """Re-apply PX patch with different settings (no weight reload).""" entry = self._models[model_id] registry = entry["registry"] print(f"[ModelManager] Re-patching {model_id} (subjective={px_subjective}, gamma={px_gamma}, routing={px_routing_mode}, preset={px_config_preset})...") # Remove existing patch remove_fn = self._get_patch_function(model_id, "remove_px_patch") if remove_fn: try: remove_fn(entry["model"]) except Exception as e: print(f"[ModelManager] Warning: remove_px_patch failed: {e}") # Re-apply with new settings (if not BASELINE) if px_config_preset != "BASELINE": patch_kwargs = dict(registry["patch_kwargs"]) if px_subjective and registry.get("subjective_kwargs"): patch_kwargs.update(registry["subjective_kwargs"]) if px_config_preset is not None: patch_kwargs["config_preset"] = _migrate_preset(px_config_preset) if px_gamma is not None: patch_kwargs["gamma"] = px_gamma if px_routing_mode is not None: patch_kwargs["routing_mode"] = px_routing_mode # SR-61b: Explicitly pass subjective flag patch_kwargs["subjective_enabled"] = px_subjective # verstärkbar Relay-Parameter (psychomotrik seite15) → patch_kwargs if px_relay_sign is not None: patch_kwargs["relay_sign"] = px_relay_sign if px_relay_alpha is not None: patch_kwargs["relay_alpha"] = px_relay_alpha if px_relay_layer is not None: patch_kwargs["relay_layer"] = px_relay_layer apply_fn = self._get_patch_function(model_id, "apply_px_patch") apply_fn(entry["model"], **patch_kwargs) tm = self._resolve_text_model(entry["model"]) entry["model"].tokenizer = tm.tokenizer = entry["tokenizer"] else: print(f"[ModelManager] {model_id} returned to baseline state.") entry["px_subjective"] = px_subjective entry["px_gamma"] = px_gamma entry["px_routing_mode"] = px_routing_mode entry["px_config_preset"] = px_config_preset entry["px_relay_sign"] = px_relay_sign entry["px_relay_alpha"] = px_relay_alpha entry["px_relay_layer"] = px_relay_layer def _get_patch_function(self, model_id: str, function_name: str): """Import patch module and get a function by name.""" registry = MODEL_REGISTRY[model_id] patch_dir = registry["patch_dir"] # Ensure px_patches directory is on sys.path base_dir = os.path.dirname(os.path.abspath(__file__)) px_dir = os.path.join(base_dir, "px_patches") if px_dir not in sys.path: sys.path.insert(0, px_dir) # Import the patch module module_name = f"{patch_dir}.patch" if module_name not in sys.modules: mod = importlib.import_module(module_name) else: mod = sys.modules[module_name] fn = getattr(mod, function_name, None) if fn is None: raise AttributeError(f"Function {function_name} not found in {module_name}") return fn def get_px_metrics(self, model_id: str) -> dict: """Get latest PX cognitive metrics for a model.""" registry = MODEL_REGISTRY.get(model_id, {}) if registry.get("patch_dir") is None: return {"patched": False, "model_type": registry.get("model_type", "unknown")} if model_id not in self._models: return {} entry = self._models[model_id] model = entry["model"] # Try get_px_metrics function (MiniCPM5 version) try: get_fn = self._get_patch_function(model_id, "get_px_metrics") return get_fn(model) except (AttributeError, ImportError): pass # Fallback: read attributes directly text_model = self._resolve_text_model(model) if text_model is None: return {} return { "phi": getattr(text_model, "_px_phi", 1.0), "steps": getattr(text_model, "_px_loops_run", 0), "path": getattr(text_model, "_px_path", []), "zone": getattr(text_model, "_px_zone", "UNKNOWN"), "zone_weights": getattr(text_model, "_px_zone_weights", {}), "cognitive_signature": getattr(text_model, "_px_cognitive_signature", {}), "subjective": getattr(text_model, "_px_subjective_enabled", False), } def _resolve_text_model(self, model): """Find the transformer backbone in the model.""" if hasattr(model, "layers") and hasattr(model, "rotary_emb"): return model for name, mod in model.named_modules(): if hasattr(mod, "layers") and hasattr(mod, "rotary_emb"): return mod return model.model if hasattr(model, "model") else model def list_models(self) -> list: """Return list of available model IDs.""" return list(MODEL_REGISTRY.keys()) def is_loaded(self, model_id: str) -> bool: """Check if a model is currently loaded.""" return model_id in self._models def is_px_model(self, model_id: str) -> bool: """Check if a model has PX patch enabled.""" registry = MODEL_REGISTRY.get(model_id, {}) return registry.get("patch_dir") is not None def unload(self, model_id: str): """Unload a model from GPU memory.""" if model_id in self._models: entry = self._models.pop(model_id) # Remove patch references before deleting model try: registry = entry.get("registry", {}) if registry.get("patch_dir") is not None: remove_fn = self._get_patch_function(model_id, "remove_px_patch") if remove_fn: remove_fn(entry["model"]) except Exception: pass # Best-effort cleanup del entry["model"] del entry["tokenizer"] if model_id in self._px_metrics: del self._px_metrics[model_id] if model_id in self._last_used: del self._last_used[model_id] if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"[ModelManager] {model_id} unloaded.") async def shutdown(self): """Cleanup all loaded models.""" for model_id in list(self._models.keys()): entry = self._models[model_id] del entry["model"] del entry["tokenizer"] self._models.clear() if torch.cuda.is_available(): torch.cuda.empty_cache() print("[ModelManager] All models unloaded.")