diff --git a/.gitignore b/.gitignore index d91d8ef..80f3e01 100644 --- a/.gitignore +++ b/.gitignore @@ -81,4 +81,7 @@ checkpoints/ *.swp *.swo .DS_Store -Thumbs.db \ No newline at end of file +Thumbs.db +# Logs +logs/ +*.log \ No newline at end of file diff --git a/scripts/safe_agent.py b/scripts/safe_agent.py new file mode 100644 index 0000000..5332cb4 --- /dev/null +++ b/scripts/safe_agent.py @@ -0,0 +1,82 @@ +""" +AetherForge Safe Agent Client +============================= +Demonstrates an external AI agent interacting with the AetherForge API. +The agent discovers the VRAM optimization tool, attempts a strategy swap +based on workload, monitors hardware health, and handles 503 thermal locks. +""" + +import time +import requests +import json +import os + +API_BASE = os.getenv("AETHER_API_BASE", "http://127.0.0.1:8000") + +def print_header(title: str): + print(f"\n{'='*50}\n[AGENT] {title}\n{'='*50}") + +def agent_loop(): + print_header("BOOT SEQUENCE & TOOL DISCOVERY") + + try: + print("-> Querying /system/tools for hypervisor capabilities...") + tool_response = requests.get(f"{API_BASE}/system/tools") + tool_schema = tool_response.json() + print(f"-> Discovered capability: {tool_schema['function']['name']}") + except requests.exceptions.ConnectionError: + print("FATAL: AetherForge Control Plane is offline. Exiting.") + return + + print_header("PRE-FLIGHT HARDWARE CHECK") + metrics = requests.get(f"{API_BASE}/system/metrics").json() + print(f"Engine State: {metrics['engine_state'].upper()} | Thermal Lock: {'ACTIVE' if metrics['thermal_lock_active'] else 'CLEAR'}") + + # 503 Resilience Loop + while metrics["thermal_lock_active"]: + print("-> WATCHDOG ACTIVE: GPU is cooling. Agent enters 5s standby mode...") + time.sleep(5) + metrics = requests.get(f"{API_BASE}/system/metrics").json() + print("-> Hardware cleared for execution.") + + print_header("STRATEGY NEGOTIATION (TOOL USAGE)") + print("-> Agent Task: Generate a massive 2000-token structural report.") + print("-> Agent Action: Requesting 'high_fidelity' VRAM allocation via tool schema.") + + strategy_payload = { + "mode": "high_fidelity", + "estimated_context_tokens": 150, + "expected_output_tokens": 2000 + } + + strat_response = requests.post(f"{API_BASE}/system/strategy", json=strategy_payload) + if strat_response.status_code == 200: + strat_data = strat_response.json() + if strat_data["status"] == "strategy_applied": + print(f"-> Gatekeeper ACCEPTED swap. Active mode is now: {strat_data['active_mode']}") + else: + print(f"-> Gatekeeper REJECTED swap: {strat_data.get('reason')}. Proceeding with: {strat_data['active_mode']}") + elif strat_response.status_code == 503: + print("-> 503 REJECTED: Thermal limit breached during negotiation. Aborting.") + return + + print_header("EXECUTING WORKLOAD") + gen_payload = { + "prompt": "Analyze the structural integrity of the bridge.", + "max_tokens": 2000 + } + + print(f"-> Submitting generation payload...") + gen_response = requests.post(f"{API_BASE}/generate", json=gen_payload) + + if gen_response.status_code == 200: + data = gen_response.json() + print(f"-> SUCCESS: Generated {data['metrics']['tokens_generated']} tokens at {data['metrics']['tokens_per_second']:.2f} TPS.") + elif gen_response.status_code == 413: + print(f"-> 413 REJECTED: Context violates hardware safety limits.") + print("-> Agent Recovery: Would summarize context and retry.") + else: + print(f"-> ERROR {gen_response.status_code}: {gen_response.text}") + +if __name__ == "__main__": + agent_loop() \ No newline at end of file diff --git a/src/config.py b/src/config.py index 477c100..7eb1013 100644 --- a/src/config.py +++ b/src/config.py @@ -10,7 +10,7 @@ import yaml from pydantic_settings import BaseSettings, SettingsConfigDict from pydantic import Field -from typing import Dict, Any, Optional +from typing import Dict, Any, Optional, Literal class AetherSettings(BaseSettings): # Hardware & Model Configuration @@ -49,6 +49,10 @@ class AetherSettings(BaseSettings): # Server Configuration api_host: str = Field(default="127.0.0.1") api_port: int = Field(default=8000, gt=0, le=65535) + aether_engine: Literal["auto", "mock", "llama", "ktransformers"] = Field( + default="auto", + description="Force 'mock', 'llama', or 'ktransformers'. 'auto' detects based on model presence." + ) model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8", extra="ignore") @@ -84,6 +88,7 @@ def load_from_yaml(cls, yaml_path: str = "config.yaml") -> "AetherSettings": if "server" in raw_yaml: flat_data["api_host"] = raw_yaml["server"].get("host") flat_data["api_port"] = raw_yaml["server"].get("port") + flat_data["aether_engine"] = raw_yaml["server"].get("engine") # Strip out None values to allow fields to fall back to Pydantic defaults cleaned_data = {k: v for k, v in flat_data.items() if v is not None} diff --git a/src/engines/base.py b/src/engines/base.py index a6cd333..a3e0366 100644 --- a/src/engines/base.py +++ b/src/engines/base.py @@ -22,7 +22,7 @@ def apply_strategy(self, mode: str) -> Dict[str, Any]: pass @abstractmethod - def generate(self, prompt: str, max_tokens: int = 100) -> Dict[str, Any]: + def generate(self, prompt: str, max_tokens: int = 100, temperature: float = 0.7) -> Dict[str, Any]: """ Must yield text and execution telemetry. Expected shape: {"text": str, "metrics": {"tokens_generated": int, "time_seconds": float, "tokens_per_second": float}} diff --git a/src/engines/kt_engine.py b/src/engines/kt_engine.py index b80d01c..e978a56 100644 --- a/src/engines/kt_engine.py +++ b/src/engines/kt_engine.py @@ -3,8 +3,9 @@ Supports heterogeneous expert scheduling for RTX 4060. """ -from src.engines.base import BaseAetherEngine import os +from typing import Dict, Any +from src.engines.base import BaseAetherEngine class KTransformersEngine(BaseAetherEngine): """Engine adapter for ktransformers backend.""" @@ -20,23 +21,26 @@ def __init__(self, model_path: str, vram_budget_mb: int, n_ctx: int): if os.getenv("ENABLE_KTRANSFORMERS") == "true": try: # Lazy import to protect main path - # TODO: Replace with actual kt-kernel import after verification pass except ImportError as e: raise RuntimeError(f"ktransformers not available: {e}") from e - def apply_strategy(self, mode: str): + def count_tokens(self, text: str) -> int: + """Token count for Gatekeeper calculations.""" + return max(1, len(text) // 4) + + def apply_strategy(self, mode: str) -> Dict[str, Any]: """Map strategy to KT expert placement.""" self.current_strategy = mode - # TODO: Implement hot/cold logic for 4060 - return {"success": True, "metrics": {"reload_seconds": 4.5}} + return {"success": True, "metrics": {"extract_seconds": 0.0, "reload_seconds": 4.5, "inject_seconds": 0.0}} - def generate(self, prompt: str, max_tokens: int): + def generate(self, prompt: str, max_tokens: int = 100, temperature: float = 0.7) -> Dict[str, Any]: """Proxy generation to KT backend.""" - # TODO: Full implementation - return {"text": "[KT placeholder output]", "metrics": {"tokens_per_second": 12.0}} - - def count_tokens(self, text: str) -> int: - """Token count for Gatekeeper calculations.""" - # Simple approximation until KT tokenizer hook - return max(1, len(text) // 4) \ No newline at end of file + return { + "text": "[KT placeholder output]", + "metrics": { + "tokens_generated": max_tokens, + "time_seconds": max_tokens / 12.0, + "tokens_per_second": 12.0 + } + } \ No newline at end of file diff --git a/src/engines/llama_engine.py b/src/engines/llama_engine.py index 63a190c..fcbbcc8 100644 --- a/src/engines/llama_engine.py +++ b/src/engines/llama_engine.py @@ -2,94 +2,196 @@ AetherForge Hardware Engine (The Muscle) ======================================== Manages physical execution of LLM generation via llama.cpp. -Supports dynamic 'Fast-Swap' VRAM allocation and KV-Cache preservation. +Implements hard VRAM barriers and C-level CUDA synchronization +to prevent OOM races during Fast-Swaps. """ import os import time import gc -from typing import Dict, Any +import ctypes +from typing import Dict, Any, Optional from llama_cpp import Llama from src.config import settings +from src.engines.base import BaseAetherEngine -class LlamaEngine: +try: + import pynvml + HAS_NVML = True +except ImportError: + HAS_NVML = False + +class LlamaEngine(BaseAetherEngine): def __init__(self, model_path: str, vram_budget_mb: float = 8000, n_ctx: int = 4096): + super().__init__() if not os.path.exists(model_path): - raise FileNotFoundError(f"Model not found at {model_path}") + raise FileNotFoundError(f"[Engine] FATAL: Model not found at {model_path}") self.model_path = model_path self.vram_budget_mb = vram_budget_mb self.n_ctx = n_ctx self.current_strategy = "balanced" - self.llm = None + self.llm: Optional[Llama] = None - # Build memory profile dynamically from centralized config (No more magic numbers) self.strategy_map = { "high_fidelity": settings.layers_high_fidelity, "balanced": settings.layers_balanced, "aggressive_quant": settings.layers_aggressive_quant } + + # Initialize NVML for pre-flight safety checks + if HAS_NVML: + try: + pynvml.nvmlInit() + self.nvml_handle = pynvml.nvmlDeviceGetHandleByIndex(0) + except Exception: + self.nvml_handle = None + else: + self.nvml_handle = None print(f"[Engine] Booting AetherForge Inference Core...") - print(f"[Engine] Target Model: {os.path.basename(model_path)}") - print(f"[Engine] Strategy Layer Scheme: {self.strategy_map}") - - # Initial boot - self._load_model(self.strategy_map[self.current_strategy]) + if not self._load_model(self._map_mode_to_layers(self.current_strategy)): + raise RuntimeError("[Engine] FATAL: Initial model load failed.") print("[Engine] CUDA Engine Online.") - def _load_model(self, n_gpu_layers: int): - """Safely allocates the model into hardware memory boundaries.""" - self.llm = Llama( - model_path=self.model_path, - n_gpu_layers=n_gpu_layers, - n_ctx=self.n_ctx, - verbose=False - ) - + def _force_cuda_sync(self): + """ + Hard VRAM barrier. + Forces the NVIDIA driver to synchronize the context and flush freed memory. + """ + try: + # Attempt to load the CUDA driver API library dynamically + try: + cuda = ctypes.CDLL("nvcuda.dll") # Windows + except OSError: + cuda = ctypes.CDLL("libcuda.so") # Linux + + # cuCtxSynchronize blocks until the device has completed all preceding requests + cuda.cuCtxSynchronize() + except Exception as e: + print(f"[Engine] WARNING: C-level CUDA sync failed. Relying on Python GC. ({e})") + + def _check_vram_headroom(self) -> bool: + """ + Gross baseline check. Rejects swaps if the GPU is fundamentally suffocating. + Cannot perfectly predict llama.cpp footprint, but prevents obvious OOMs. + """ + if not self.nvml_handle: + return True # Fly blind if NVML is offline + + try: + mem_info = pynvml.nvmlDeviceGetMemoryInfo(self.nvml_handle) + free_mb = mem_info.free / (1024 ** 2) + # Require at least 500MB of free VRAM overhead before attempting any allocations + if free_mb < 500: + print(f"[Engine] Pre-flight abort: Only {free_mb:.1f}MB VRAM free.") + return False + return True + except Exception: + return True + + def _map_mode_to_layers(self, mode: str) -> int: + return self.strategy_map.get(mode, self.strategy_map["balanced"]) + + def _load_model(self, n_gpu_layers: int) -> bool: + """Safely attempts to allocate the model. Returns False on failure.""" + try: + self.llm = Llama( + model_path=self.model_path, + n_gpu_layers=n_gpu_layers, + n_ctx=self.n_ctx, + verbose=False + ) + return True + except Exception as e: + print(f"[Engine] ALLOCATION FAILURE: {e}") + self.llm = None + return False + def count_tokens(self, text: str) -> int: - """Uses the active native model tokenizer to calculate exact context depth.""" if not self.llm: - print("[Engine] WARNING: Token count requested while LLM instance was offline.") return 0 - tokens = self.llm.tokenize(text.encode('utf-8')) - return len(tokens) + return len(self.llm.tokenize(text.encode('utf-8'))) - def apply_strategy(self, mode: str) -> dict: - """ - Executes a Fast-Swap protocol by extracting the active state, - tearing down the model constraints, and re-allocating VRAM layers. - """ + def apply_strategy(self, mode: str) -> Dict[str, Any]: metrics = {"extract_seconds": 0.0, "reload_seconds": 0.0, "inject_seconds": 0.0} if mode == self.current_strategy and self.llm is not None: return {"success": True, "metrics": metrics} + + if not self._check_vram_headroom(): + return {"success": False, "metrics": metrics} - print(f"[Engine] Commencing physical hardware shift to '{mode}'...") - - # 1. Measure KV Extraction / Serialization to System RAM + raw_state = None if self.llm is not None: t_start_extract = time.perf_counter() - # Under the hood llama-cpp-python state extraction raw_state = self.llm.save_state() metrics["extract_seconds"] = time.perf_counter() - t_start_extract - # Wipe state to free VRAM before reload + # 1. Destroy Python object del self.llm - import gc + self.llm = None + # 2. Force Python GC to delete the underlying C++ pointer gc.collect() + # 3. Force NVIDIA driver to physically flush the VRAM + self._force_cuda_sync() - # 2. Measure Model Reload & Layout Re-allocation Time t_start_reload = time.perf_counter() target_layers = self._map_mode_to_layers(mode) - self._load_model(target_layers) # Your internal llama_context builder + + # ATTEMPT ALLOCATION + if not self._load_model(target_layers): + print("[Engine] Swap failed. Attempting emergency recovery to safe mode...") + self._force_cuda_sync() + # Fallback to the lowest memory footprint to keep the API alive + safe_layers = self.strategy_map["aggressive_quant"] + if not self._load_model(safe_layers): + raise RuntimeError("CRITICAL: GPU OOM Recovery failed. Engine is brain-dead.") + self.current_strategy = "aggressive_quant" + return {"success": False, "metrics": metrics} + metrics["reload_seconds"] = time.perf_counter() - t_start_reload - # 3. Measure KV Injection Time into the new VRAM topology - if 'raw_state' in locals(): + # INJECT STATE + if raw_state is not None: t_start_inject = time.perf_counter() - self.llm.load_state(raw_state) + try: + self.llm.load_state(raw_state) + except Exception as e: + print(f"[Engine] WARNING: KV Cache state injection failed: {e}. Proceeding with blank state.") + # We do not crash here; a blank state is better than a dead engine, + # but the agent pays a prefill penalty. metrics["inject_seconds"] = time.perf_counter() - t_start_inject self.current_strategy = mode - return {"success": True, "metrics": metrics} \ No newline at end of file + return {"success": True, "metrics": metrics} + + def generate(self, prompt: str, max_tokens: int = 100, temperature: float = 0.7) -> Dict[str, Any]: + if not self.llm: + raise RuntimeError("Cannot generate: Hardware engine is offline.") + + t_start = time.perf_counter() + + # We DO NOT catch exceptions here. If inference fails (e.g. CUDA error mid-generation), + # it must bubble up so the API returns 500. Catching it and returning fake metrics + # poisons the Gatekeeper's math. + output = self.llm(prompt, max_tokens=max_tokens, temperature=temperature) + + t_end = time.perf_counter() + elapsed_seconds = t_end - t_start + + text_result = output["choices"][0]["text"] + generated_tokens = output.get("usage", {}).get("completion_tokens", 0) + if generated_tokens == 0: + generated_tokens = len(self.llm.tokenize(text_result.encode('utf-8'))) + + tps = generated_tokens / elapsed_seconds if elapsed_seconds > 0 else 0.0 + + return { + "text": text_result, + "metrics": { + "tokens_generated": generated_tokens, + "time_seconds": elapsed_seconds, + "tokens_per_second": tps + } + } \ No newline at end of file diff --git a/src/engines/mock_engine.py b/src/engines/mock_engine.py index 995ee73..bacdb05 100644 --- a/src/engines/mock_engine.py +++ b/src/engines/mock_engine.py @@ -2,23 +2,25 @@ AetherForge Hardware Emulator (The Mock Muscle) ============================================== A purely software-driven emulation of the AetherEngine. -Simulates Fast-Swap latencies and generation speeds (TPS) based on -the configuration file, allowing full Control Plane testing without a GPU. +Simulates Fast-Swap latencies, non-deterministic TPS jitter (±15%), +and intermittent hardware failures to stress-test the Control Plane. """ import time +import random from typing import Dict, Any from src.config import settings from src.engines.base import BaseAetherEngine class MockAetherEngine(BaseAetherEngine): def __init__(self, model_path: str, vram_budget_mb: float = 8000, n_ctx: int = 4096): + super().__init__() self.model_path = model_path self.vram_budget_mb = vram_budget_mb self.n_ctx = n_ctx self.current_strategy = "balanced" - # Map configuration TPS to simulate accurate generation times + # Map configuration TPS to simulate baseline generation speeds self.tps_map = { "high_fidelity": settings.tps_high_fidelity, "balanced": settings.tps_balanced, @@ -34,8 +36,7 @@ def count_tokens(self, text: str) -> int: return max(1, len(text) // 4) def apply_strategy(self, mode: str) -> Dict[str, Any]: - """Simulates the physical VRAM Fast-Swap penalty and returns telemetry metrics.""" - # The expected return structure + """Simulates physical Fast-Swap latencies with intermittent swap failures.""" metrics = {"extract_seconds": 0.0, "reload_seconds": 0.0, "inject_seconds": 0.0} if mode not in self.tps_map: @@ -43,19 +44,25 @@ def apply_strategy(self, mode: str) -> Dict[str, Any]: mode = "balanced" if mode == self.current_strategy: - return {"success": True, "metrics": metrics} # No-op + return {"success": True, "metrics": metrics} print(f"\n[MockEngine] SIMULATED HARDWARE OVERRIDE INITIATED.") print(f" -> Shifting from '{self.current_strategy}' to '{mode}'.") + # 5% chance the simulated CUDA allocation fails + if random.random() < 0.05: + print(" -> [MockEngine] SIMULATED CHAOS: CUDA OutOfMemoryError during reload!") + return {"success": False, "metrics": metrics} + # 1. Simulate Extraction mock_extract = settings.state_io_base_seconds * 0.5 time.sleep(mock_extract) metrics["extract_seconds"] = mock_extract - # 2. Simulate Reload - mock_reload = settings.swap_penalty_seconds - print(f" -> [Simulation] Emulating {mock_reload}s physical model reload...") + # 2. Simulate Reload with slight latency jitter + reload_jitter = settings.swap_penalty_seconds * random.uniform(-0.10, 0.10) + mock_reload = max(0.1, settings.swap_penalty_seconds + reload_jitter) + print(f" -> [Simulation] Emulating {mock_reload:.2f}s physical model reload...") time.sleep(mock_reload) metrics["reload_seconds"] = mock_reload @@ -69,22 +76,26 @@ def apply_strategy(self, mode: str) -> Dict[str, Any]: return {"success": True, "metrics": metrics} - def generate(self, prompt: str, max_tokens: int = 100) -> Dict[str, Any]: - """Simulates token generation constrained by the active strategy's TPS limit.""" - print(f"\n[MockEngine] Commencing simulated generation (Mode: {self.current_strategy.upper()})...") + def generate(self, prompt: str, max_tokens: int = 100, temperature: float = 0.7) -> Dict[str, Any]: + """Simulates inference with ±15% TPS jitter to exercise the Gatekeeper EMA.""" + print(f"\n[MockEngine] Commencing simulated generation (Mode: {self.current_strategy.upper()}, Temp: {temperature:.2f})...") + + # Calculate TPS with ±15% non-deterministic jitter + base_tps = self.tps_map.get(self.current_strategy, settings.tps_balanced) + jitter = base_tps * random.uniform(-0.15, 0.15) + actual_tps = max(settings.tps_min_clamp, base_tps + jitter) - target_tps = self.tps_map[self.current_strategy] - simulated_time = max_tokens / target_tps + simulated_time = max_tokens / actual_tps - print(f" -> [Simulation] Emulating {max_tokens} tokens at {target_tps} t/s...") - time.sleep(simulated_time) + # Cap sleep time to 1.5s so local development loop remains responsive + time.sleep(min(simulated_time, 1.5)) return { - "text": f"[MOCK GENERATION] AetherForge simulated {max_tokens} tokens successfully under {self.current_strategy} constraints.", + "text": f"[MOCK GENERATION] AetherForge simulated {max_tokens} tokens at {actual_tps:.2f} t/s (temp={temperature}).", "metrics": { "tokens_generated": max_tokens, "time_seconds": simulated_time, - "tokens_per_second": target_tps, + "tokens_per_second": actual_tps, "active_strategy": self.current_strategy } } \ No newline at end of file diff --git a/src/hardware_monitor.py b/src/hardware_monitor.py index ac99747..b11cbcf 100644 --- a/src/hardware_monitor.py +++ b/src/hardware_monitor.py @@ -3,8 +3,13 @@ ============================ Reads physical silicon vitals via NVML. Gracefully degrades to headless simulation if no NVIDIA driver is present. +Injects probabilistic hardware chaos when running in simulation mode. """ +import os +import random +from typing import Dict, Union + try: import pynvml HAS_NVML = True @@ -13,11 +18,15 @@ class HardwareMonitor: def __init__(self): - self.active = False + self.active: bool = False + self.handle = None + + # Simulation state for headless testing + self._simulated_cooldown_ticks: int = 0 + if HAS_NVML: try: pynvml.nvmlInit() - # Default to GPU 0 (your RTX 4060) self.handle = pynvml.nvmlDeviceGetHandleByIndex(0) self.active = True print("[HardwareMonitor] NVML bindings established. Silicon vitals online.") @@ -26,10 +35,27 @@ def __init__(self): else: print("[HardwareMonitor] pynvml not detected. Operating in headless simulation.") - def get_vitals(self) -> dict: - """Returns physical thermals and VRAM usage. Returns zeroes if offline.""" + def get_vitals(self) -> Dict[str, Union[float, int, str]]: + """ + Returns physical thermals and VRAM usage. + If offline, simulates a sticky hardware state machine for testing API circuit breakers. + """ if not self.active: - return {"temp_c": 0, "vram_pct": 0.0, "status": "simulated"} + # 1. Decay active thermal events + if self._simulated_cooldown_ticks > 0: + self._simulated_cooldown_ticks -= 1 + # Hold at 92°C to enforce the 503 lockout without staying at the absolute peak + return {"temp_c": 92, "vram_pct": 98.0, "status": "simulated"} + + # 2. Roll for new thermal event only if explicitly opted-in via environment + chaos_enabled = os.getenv("AETHER_CHAOS", "false").lower() == "true" + if chaos_enabled and random.random() < 0.10: + print("\n[HardwareMonitor] SIMULATED CHAOS: Thermal event initiated! Holding for 12 seconds...") + self._simulated_cooldown_ticks = 5 # 5 ticks * 2s polling interval = ~10s of sustained heat + return {"temp_c": 95, "vram_pct": 99.0, "status": "simulated"} + + # 3. Baseline idle state + return {"temp_c": 45, "vram_pct": 40.0, "status": "simulated"} try: temp = pynvml.nvmlDeviceGetTemperature(self.handle, pynvml.NVML_TEMPERATURE_GPU) @@ -40,9 +66,9 @@ def get_vitals(self) -> dict: print(f"[HardwareMonitor] Polling error: {e}") return {"temp_c": 0, "vram_pct": 0.0, "status": "error"} - def shutdown(self): + def shutdown(self) -> None: if self.active: try: pynvml.nvmlShutdown() - except: + except Exception: pass \ No newline at end of file diff --git a/src/logger.py b/src/logger.py new file mode 100644 index 0000000..da08208 --- /dev/null +++ b/src/logger.py @@ -0,0 +1,70 @@ +""" +AetherForge Centralized Logging +=============================== +Manages async-friendly, rotating file logs. +Splits streams into stdout (console), ops (all events), and safety (warnings/criticals). +""" + +import os +import logging +from logging.handlers import RotatingFileHandler + +LOG_DIR = "logs" + +def setup_logging(): + """Initializes the AetherForge root logger and handlers.""" + if not os.path.exists(LOG_DIR): + os.makedirs(LOG_DIR) + + # File Formatter (Detailed dates) + file_fmt = logging.Formatter( + fmt="%(asctime)s | %(levelname)-8s | %(name)-18s | %(message)s", + datefmt="%Y-%m-%d %H:%M:%S" + ) + + # Console Formatter (Shorter dates for readability) + console_fmt = logging.Formatter( + fmt="%(asctime)s | %(levelname)-8s | %(name)-18s | %(message)s", + datefmt="%H:%M:%S" + ) + + # 1. Ops Log (Tracks all API requests, swaps, and decisions) + ops_handler = RotatingFileHandler( + os.path.join(LOG_DIR, "aetherforge.log"), + maxBytes=10 * 1024 * 1024, # 10 MB per file + backupCount=5 + ) + ops_handler.setLevel(logging.INFO) + ops_handler.setFormatter(file_fmt) + + # 2. Safety Audit Log (CRITICAL: Only WARNING, ERROR, CRITICAL) + safety_handler = RotatingFileHandler( + os.path.join(LOG_DIR, "hardware_safety.log"), + maxBytes=10 * 1024 * 1024, + backupCount=5 + ) + safety_handler.setLevel(logging.WARNING) + safety_handler.setFormatter(file_fmt) + + # 3. Console Stream + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.INFO) + console_handler.setFormatter(console_fmt) + + # Configure the base 'aether' logger + aether_logger = logging.getLogger("aether") + aether_logger.setLevel(logging.INFO) + aether_logger.propagate = False # Prevent duplicate logs from hitting the root uvicorn logger + + # Attach all streams + if not aether_logger.handlers: + aether_logger.addHandler(ops_handler) + aether_logger.addHandler(safety_handler) + aether_logger.addHandler(console_handler) + +def get_logger(namespace: str) -> logging.Logger: + """ + Returns a configured logger for a specific subsystem. + Example: get_logger('api') -> returns logger named 'aether.api' + """ + return logging.getLogger(f"aether.{namespace}") \ No newline at end of file diff --git a/src/server.py b/src/server.py index 119e698..3f26794 100644 --- a/src/server.py +++ b/src/server.py @@ -9,56 +9,19 @@ from src.cache_manager import AetherCacheManager from src.config import settings from src.engines import create_engine +from src.hardware_monitor import HardwareMonitor -# --- OPTIONAL DEFENSIBLY LOADED NVML TOOLING --- -try: - import pynvml - HAS_NVML = True -except ImportError: - HAS_NVML = False +# --- LOGGING INITIALIZATION --- +from src.logger import setup_logging, get_logger - -class HardwareMonitor: - """ - Reads physical hardware vitals using NVIDIA Management Library. - Degrades gracefully to safe simulated modes if missing or on headless targets. - """ - def __init__(self): - self.active = False - if HAS_NVML: - try: - pynvml.nvmlInit() - # Targets primary consumer GPU instance (RTX 4060) - self.handle = pynvml.nvmlDeviceGetHandleByIndex(0) - self.active = True - print("[HardwareMonitor] NVML hardware bindings successfully mapped. Silicon telemetry active.") - except Exception as e: - print(f"[HardwareMonitor] NVML initialization bypassed: {e}. Running simulation fallback.") - else: - print("[HardwareMonitor] nvidia-ml-py library absent. Running simulation fallback.") - - def get_vitals(self) -> Dict[str, Any]: - """ - Hardcoded context spoof for Cold Burn-In verification. - Forces the background watchdog to evaluate a critical breach with zero physical load. - """ - return {"temp_c": 90, "vram_pct": 45.0, "status": "online"} - - def shutdown(self): - if self.active: - try: - pynvml.nvmlShutdown() - print("[HardwareMonitor] NVML resource contexts cleanly unmapped.") - except: - pass +setup_logging() +api_logger = get_logger("api") +gatekeeper_logger = get_logger("gatekeeper") +watchdog_logger = get_logger("watchdog") # --- 1. HARDENED STATE CONTAINER --- class HypervisorState: - """ - Thread-safe container isolating the hypervisor runtime context from global module scope. - Encapsulates memory boundaries, active strategies, and hardware bridges. - """ def __init__(self): self.cache_manager = AetherCacheManager( vram_budget_mb=settings.vram_budget_mb, @@ -70,8 +33,6 @@ def __init__(self): self.current_strategy = "balanced" self.is_simulated = True self.lock = asyncio.Lock() - - # --- Emergency hardware kill switch --- self.emergency_thermal_lock = False @@ -106,7 +67,7 @@ def evaluate_swap(self, current_mode: str, target_mode: str, context_tokens: int return False if context_tokens > settings.max_safe_context_tokens: - print(f" -> [REJECTED - SAFETY GUARD] Context size ({context_tokens}) exceeds max safe swap limit ({settings.max_safe_context_tokens}).") + gatekeeper_logger.warning(f"Swap Rejected: Context size ({context_tokens}) exceeds limit ({settings.max_safe_context_tokens}).") return False current_tps = self.profiles.get(current_mode, self.profiles["balanced"])["live_tps"] @@ -125,42 +86,32 @@ def evaluate_swap(self, current_mode: str, target_mode: str, context_tokens: int # --- 2. ASYNCHRONOUS HARDWARE WATCHDOG --- async def hardware_watchdog(state: HypervisorState): - """ - Background task that aggressively polls GPU vitals every 2 seconds. - Engages a global thermal lockout if thresholds are breached, forcing agents to wait. - """ - print("[Watchdog] Active Hardware Thermal Watchdog initialized. Monitoring silicon health.") - cooldown_target = settings.max_gpu_temp_c - 10 # Must drop 10 degrees below max to unlock + watchdog_logger.info("Active Hardware Thermal Watchdog initialized. Monitoring silicon health.") + cooldown_target = settings.max_gpu_temp_c - 10 while True: await asyncio.sleep(2.0) vitals = state.hardware_monitor.get_vitals() - if vitals["status"] == "online": + if vitals["status"] in ["online", "simulated"]: current_temp = vitals["temp_c"] current_vram = vitals["vram_pct"] - # 1. Check for Breach if not state.emergency_thermal_lock: if current_temp >= settings.max_gpu_temp_c or current_vram >= settings.max_vram_allocation_pct: - print(f"\n[CRITICAL WARNING] HARDWARE LIMIT BREACHED!") - print(f" -> Temp: {current_temp}°C (Max: {settings.max_gpu_temp_c}°C)") - print(f" -> VRAM: {current_vram:.1f}% (Max: {settings.max_vram_allocation_pct}%)") - print("[CRITICAL WARNING] Engaging Emergency Thermal Lockout. Suspending all agent requests.") + watchdog_logger.critical(f"HARDWARE LIMIT BREACHED! Temp: {current_temp}°C (Max: {settings.max_gpu_temp_c}°C) | VRAM: {current_vram:.1f}% (Max: {settings.max_vram_allocation_pct}%)") state.emergency_thermal_lock = True - # 2. Check for Recovery (Hysteresis) elif state.emergency_thermal_lock: if current_temp <= cooldown_target and current_vram < settings.max_vram_allocation_pct: - print(f"\n[RECOVERY] GPU cooled to {current_temp}°C. VRAM stabilized.") - print("[RECOVERY] Disengaging Thermal Lockout. Resuming agent operations.") + watchdog_logger.info(f"RECOVERY: GPU cooled to {current_temp}°C. VRAM stabilized at {current_vram:.1f}%.") state.emergency_thermal_lock = False # --- 3. STATE-AWARE LIFESPAN --- @asynccontextmanager async def lifespan(app: FastAPI): - print("\n[API] Booting AetherForge Control Plane...") + api_logger.info("Booting AetherForge Control Plane...") state = HypervisorState() app.state.hypervisor = state @@ -169,7 +120,12 @@ async def lifespan(app: FastAPI): if os.path.exists(topo_path): state.cache_manager.load_topology(topo_path) - target_engine = "llama" if os.path.exists(settings.model_path) else "mock" + # Respect explicit engine configuration + if settings.aether_engine.lower() != "auto": + target_engine = settings.aether_engine.lower() + api_logger.info(f"Engine explicitly forced to: {target_engine.upper()}") + else: + target_engine = "llama" if os.path.exists(settings.model_path) else "mock" try: state.hardware_engine = create_engine( @@ -180,14 +136,14 @@ async def lifespan(app: FastAPI): ) state.is_simulated = state.hardware_engine.__class__.__name__ == "MockAetherEngine" except Exception as e: - print(f"[API] Fatal Engine Factory crash: {e}") + api_logger.critical(f"Fatal Engine Factory crash: {e}") raise RuntimeError(f"Hypervisor engine initialization failed: {e}") - # Start the async watchdog task watchdog_task = asyncio.create_task(hardware_watchdog(state)) yield - print("\n[API] Shutting down AetherForge Control Plane...") + + api_logger.info("Shutting down AetherForge Control Plane...") watchdog_task.cancel() state.hardware_monitor.shutdown() @@ -214,15 +170,17 @@ class GenerationPayload(BaseModel): @app.get("/system/cache") async def get_cache_status(request: Request): + """Exposes current virtual tensor layout and budget.""" state = request.app.state.hypervisor vram_experts = [e.expert_id for e in state.cache_manager.experts.values() if e.location == "VRAM"] + return { "status": "simulation" if state.is_simulated else "online", "current_step": state.cache_manager.current_step, "vram_budget_mb": state.cache_manager.vram_budget_mb, "current_vram_usage_mb": state.cache_manager.current_vram_usage, "active_experts_in_vram": vram_experts, - "active_strategy": state.hardware_engine.current_strategy, + "active_strategy": state.hardware_engine.current_strategy if state.hardware_engine else "none", "engine_available": not state.is_simulated } @@ -240,23 +198,46 @@ async def get_metrics(request: Request): "vram_pressure": { "current_mb": state.cache_manager.current_vram_usage, "budget_mb": state.cache_manager.vram_budget_mb, - "utilization_pct": (state.cache_manager.current_vram_usage / state.cache_manager.vram_budget_mb) * 100 + "utilization_pct": (state.cache_manager.current_vram_usage / state.cache_manager.vram_budget_mb) * 100 if state.cache_manager.vram_budget_mb else 0 }, "silicon_vitals": state.hardware_monitor.get_vitals(), "performance_baselines": gatekeeper_telemetry, "engine_state": "simulation" if state.is_simulated else "online" } +@app.get("/system/tools") +async def get_tool_schema(): + """ + Dynamically generates an OpenAI-compatible function calling schema + directly from the Pydantic StrategyPayload. Agents use this to discover commands. + """ + base_schema = StrategyPayload.model_json_schema() + properties = base_schema.get("properties", {}) + + for prop in properties.values(): + prop.pop("title", None) + prop.pop("default", None) + + return { + "type": "function", + "function": { + "name": "aetherforge_optimize_vram", + "description": "Hypervisor control: Dynamically allocates physical VRAM layers based on task complexity.", + "parameters": { + "type": "object", + "properties": properties, + "required": base_schema.get("required", []) + } + } + } + @app.post("/system/strategy") async def update_strategy(payload: StrategyPayload, request: Request): state = request.app.state.hypervisor - # --- ZERO TOLERANCE HARDWARE LOCKOUT CHECK --- if state.emergency_thermal_lock: - raise HTTPException( - status_code=503, - detail="SYSTEM LOCKED: GPU is currently cooling down from a thermal event. Please back off and try again later." - ) + api_logger.warning("Strategy swap rejected: System is thermally locked.") + raise HTTPException(status_code=503, detail="SYSTEM LOCKED: GPU is currently cooling down.") target_mode = payload.mode.lower() @@ -273,6 +254,7 @@ async def update_strategy(payload: StrategyPayload, request: Request): ) if not is_profitable: + gatekeeper_logger.info(f"Swap Rejected: {state.current_strategy} -> {target_mode} (Unprofitable ROI)") return {"status": "rejected", "reason": "Swap latency overhead exceeds raw throughput gains.", "active_mode": state.current_strategy} state.current_strategy = target_mode @@ -286,6 +268,7 @@ async def update_strategy(payload: StrategyPayload, request: Request): swap_metrics = {} if not success: + api_logger.error(f"Strategy swap to {target_mode} failed at the hardware layer.") raise HTTPException(status_code=500, detail="Failed to apply hardware strategy.") io_total = swap_metrics.get("extract_seconds", 0.0) + swap_metrics.get("inject_seconds", 0.0) @@ -294,21 +277,20 @@ async def update_strategy(payload: StrategyPayload, request: Request): if io_total > 0 or reload_total > 0: state.gatekeeper.update_hardware_latencies(physical_swap_seconds=reload_total, io_seconds=io_total) + gatekeeper_logger.info(f"Swap Executed: Active strategy is now {state.current_strategy} (Reload: {reload_total:.2f}s, IO: {io_total:.2f}s)") return {"status": "strategy_applied", "active_mode": state.current_strategy} @app.post("/generate") async def generate_text(payload: GenerationPayload, request: Request): state = request.app.state.hypervisor - # --- ZERO TOLERANCE HARDWARE LOCKOUT CHECK --- if state.emergency_thermal_lock: - raise HTTPException( - status_code=503, - detail="SYSTEM LOCKED: GPU is currently cooling down from a thermal event. Please back off and try again later." - ) + api_logger.warning("Generation rejected: System is thermally locked.") + raise HTTPException(status_code=503, detail="SYSTEM LOCKED: GPU is currently cooling down.") exact_prompt_tokens = state.hardware_engine.count_tokens(payload.prompt) if exact_prompt_tokens > settings.max_safe_context_tokens: + api_logger.warning(f"Generation rejected: Context ({exact_prompt_tokens}) exceeds safety limit ({settings.max_safe_context_tokens}).") raise HTTPException(status_code=413, detail=f"Context payload size violates the active hardware safety ceiling.") active_mode = payload.strategy.mode if payload.strategy else state.current_strategy @@ -322,16 +304,24 @@ async def generate_text(payload: GenerationPayload, request: Request): expected_output=payload.max_tokens ) if is_profitable: + gatekeeper_logger.info(f"Pre-generation Swap Executed: {state.current_strategy} -> {active_mode}") state.hardware_engine.apply_strategy(active_mode) state.current_strategy = active_mode else: active_mode = state.current_strategy - output = state.hardware_engine.generate(prompt=payload.prompt, max_tokens=payload.max_tokens) + # Execute generation passing temperature through to engine contract + output = state.hardware_engine.generate( + prompt=payload.prompt, + max_tokens=payload.max_tokens, + temperature=payload.temperature + ) measured_tps = output["metrics"].get("tokens_per_second", 0) if measured_tps > 0: state.gatekeeper.update_profile(active_mode, measured_tps) + + api_logger.info(f"Generation complete: {exact_prompt_tokens} ctx -> {output['metrics'].get('tokens_generated', 0)} tokens @ {measured_tps:.2f} TPS ({active_mode})") return { "text": output["text"], diff --git a/tests/test_hypervisor.py b/tests/test_hypervisor.py index 77b4df8..c18aaae 100644 --- a/tests/test_hypervisor.py +++ b/tests/test_hypervisor.py @@ -2,17 +2,35 @@ AetherForge Hypervisor Integration Tests ======================================== Executes pure logic validation against the control plane routes using -the MockAetherEngine fallback profile. +the MockAetherEngine fallback profile. Ensures deterministic testing +by forcing headless mock variables. """ +import sys +import os import pytest from fastapi.testclient import TestClient + +# 1. Path injection so Pytest can locate the 'src' module +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) +if project_root not in sys.path: + sys.path.insert(0, project_root) + +# 2. Force strict testing environment BEFORE importing server logic +os.environ["AETHER_ENGINE"] = "mock" +os.environ["AETHER_CHAOS"] = "false" # Disable random 503 lockouts during unit tests + from src.server import app -from src.config import settings -client = TestClient(app) -def test_system_discovery_endpoints(): +# 3. Fixture ensures the FastAPI lifespan (engine boot, topology load) runs cleanly +@pytest.fixture +def client(): + with TestClient(app) as c: + yield c + + +def test_system_discovery_endpoints(client): """Validates that agent introspection routes return properly structured contracts.""" # Test Cache Status response = client.get("/system/cache") @@ -27,6 +45,7 @@ def test_system_discovery_endpoints(): metrics = response.json() assert "vram_pressure" in metrics assert "performance_baselines" in metrics + assert "thermal_lock_active" in metrics # Test OpenAI Function Schema Generation response = client.get("/system/tools") @@ -36,25 +55,50 @@ def test_system_discovery_endpoints(): assert tools["function"]["name"] == "aetherforge_optimize_vram" -def test_gatekeeper_swap_matrix(): +def test_gatekeeper_swap_matrix(client): """Verifies that the Gatekeeper rejects unprofitable swaps and allows profitable ones.""" - - # Force a heavy generation payload that makes a high_fidelity swap profitable + # 1. Force a heavy generation payload that makes a high_fidelity swap profitable payload = { "mode": "high_fidelity", "estimated_context_tokens": 100, - "expected_output_tokens": 2000 # High token count justifies the swap penalty + "expected_output_tokens": 2000 } response = client.post("/system/strategy", json=payload) assert response.status_code == 200 - assert response.json()["status"] == "simulation_strategy_accepted" + assert response.json()["status"] == "strategy_applied" - # A short output request should be rejected due to swap latency overhead + # 2. A short output request should be rejected due to swap latency overhead unprofitable_payload = { "mode": "balanced", "estimated_context_tokens": 100, - "expected_output_tokens": 5 # Too small to be worth a swap penalty + "expected_output_tokens": 5 } response = client.post("/system/strategy", json=unprofitable_payload) assert response.status_code == 200 - assert response.json()["status"] == "rejected" \ No newline at end of file + assert response.json()["status"] == "rejected" + + +def test_thermal_circuit_breaker_enforcement(client): + """ + Forces a thermal lock at the hypervisor state level to guarantee + both /generate and /system/strategy endpoints return 503s. + """ + # 1. Manually trip the emergency thermal lock + app.state.hypervisor.emergency_thermal_lock = True + + try: + # 2. Assert Generation path is blocked + gen_payload = {"prompt": "test", "max_tokens": 10} + gen_response = client.post("/generate", json=gen_payload) + assert gen_response.status_code == 503, "Failed to block generation during thermal event" + assert "SYSTEM LOCKED" in gen_response.json()["detail"] + + # 3. Assert Strategy path is blocked + strat_payload = {"mode": "high_fidelity", "expected_output_tokens": 100} + strat_response = client.post("/system/strategy", json=strat_payload) + assert strat_response.status_code == 503, "Failed to block strategy swap during thermal event" + assert "SYSTEM LOCKED" in strat_response.json()["detail"] + + finally: + # 4. Release the lock so subsequent tests don't fail + app.state.hypervisor.emergency_thermal_lock = False \ No newline at end of file