from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles from typing import List, Dict, Any, Optional import asyncio import json import time import os from dataclasses import dataclass, asdict from enum import Enum import uuid import aiohttp import logging from collections import defaultdict, deque from datetime import datetime # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Analytics and monitoring class SystemAnalytics: def __init__(self): self.start_time = time.time() self.total_connections = 0 self.active_connections = 0 self.total_interactions = 0 self.ip_addresses = defaultdict(int) # IP -> connection count self.session_durations = deque(maxlen=100) # Last 100 session durations self.waiting_times = deque(maxlen=100) # Last 100 waiting times self.users_bypassed_queue = 0 # Users who got GPU immediately self.users_waited_in_queue = 0 # Users who had to wait self.gpu_utilization_samples = deque(maxlen=100) # GPU utilization over time self.queue_size_samples = deque(maxlen=100) # Queue size over time # File handles for different analytics self.log_file = None self.gpu_metrics_file = None self.connection_events_file = None self.queue_metrics_file = None self.ip_stats_file = None self._init_log_files() def _init_log_files(self): """Initialize all analytics log files""" timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") # Main human-readable log self.log_file = f"system_analytics_{timestamp}.log" self._write_log("="*80) self._write_log("NEURAL OS MULTI-GPU SYSTEM ANALYTICS") self._write_log("="*80) self._write_log(f"System started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") self._write_log("") # Structured data files for analysis self.gpu_metrics_file = f"gpu_metrics_{timestamp}.jsonl" self.connection_events_file = f"connection_events_{timestamp}.jsonl" self.queue_metrics_file = f"queue_metrics_{timestamp}.jsonl" self.ip_stats_file = f"ip_stats_{timestamp}.jsonl" # Initialize with headers/metadata self._write_json_log(self.gpu_metrics_file, { "type": "metadata", "timestamp": time.time(), "description": "GPU utilization metrics", "fields": ["timestamp", "total_gpus", "active_gpus", "available_gpus", "utilization_percent"] }) self._write_json_log(self.connection_events_file, { "type": "metadata", "timestamp": time.time(), "description": "Connection lifecycle events", "fields": ["timestamp", "event_type", "client_id", "ip_address", "duration", "interactions", "reason"] }) self._write_json_log(self.queue_metrics_file, { "type": "metadata", "timestamp": time.time(), "description": "Queue performance metrics", "fields": ["timestamp", "queue_size", "estimated_wait", "bypass_rate", "avg_wait_time"] }) self._write_json_log(self.ip_stats_file, { "type": "metadata", "timestamp": time.time(), "description": "IP address usage statistics", "fields": ["timestamp", "ip_address", "connection_count", "total_unique_ips"] }) def _write_log(self, message): """Write message to main log file and console""" timestamp = datetime.now().strftime("%H:%M:%S") log_message = f"[{timestamp}] {message}" print(log_message) with open(self.log_file, "a") as f: f.write(log_message + "\n") def _write_json_log(self, filename, data): """Write structured data to JSON lines file""" with open(filename, "a") as f: f.write(json.dumps(data) + "\n") def log_new_connection(self, client_id: str, ip: str): """Log new connection""" self.total_connections += 1 self.active_connections += 1 self.ip_addresses[ip] += 1 unique_ips = len(self.ip_addresses) timestamp = time.time() # Human-readable log self._write_log(f"🔗 NEW CONNECTION: {client_id} from {ip}") self._write_log(f" 📊 Total connections: {self.total_connections} | Active: {self.active_connections} | Unique IPs: {unique_ips}") # Structured data logs self._write_json_log(self.connection_events_file, { "type": "connection_open", "timestamp": timestamp, "client_id": client_id, "ip_address": ip, "total_connections": self.total_connections, "active_connections": self.active_connections, "unique_ips": unique_ips }) self._write_json_log(self.ip_stats_file, { "type": "ip_update", "timestamp": timestamp, "ip_address": ip, "connection_count": self.ip_addresses[ip], "total_unique_ips": unique_ips }) def log_connection_closed(self, client_id: str, duration: float, interactions: int, reason: str = "normal"): """Log connection closed""" self.active_connections -= 1 self.total_interactions += interactions self.session_durations.append(duration) avg_duration = sum(self.session_durations) / len(self.session_durations) if self.session_durations else 0 timestamp = time.time() # Human-readable log self._write_log(f"🚪 CONNECTION CLOSED: {client_id}") self._write_log(f" ⏱️ Duration: {duration:.1f}s | Interactions: {interactions} | Reason: {reason}") self._write_log(f" 📊 Active connections: {self.active_connections} | Avg session duration: {avg_duration:.1f}s") # Structured data log self._write_json_log(self.connection_events_file, { "type": "connection_close", "timestamp": timestamp, "client_id": client_id, "duration": duration, "interactions": interactions, "reason": reason, "active_connections": self.active_connections, "avg_session_duration": avg_duration }) def log_queue_bypass(self, client_id: str): """Log when user bypasses queue (gets GPU immediately)""" self.users_bypassed_queue += 1 bypass_rate = (self.users_bypassed_queue / self.total_connections) * 100 if self.total_connections > 0 else 0 timestamp = time.time() # Human-readable log self._write_log(f"⚡ QUEUE BYPASS: {client_id} got GPU immediately") self._write_log(f" 📊 Bypass rate: {bypass_rate:.1f}% ({self.users_bypassed_queue}/{self.total_connections})") # Structured data log self._write_json_log(self.queue_metrics_file, { "type": "queue_bypass", "timestamp": timestamp, "client_id": client_id, "bypass_rate": bypass_rate, "users_bypassed": self.users_bypassed_queue, "total_connections": self.total_connections }) def log_queue_wait(self, client_id: str, wait_time: float, queue_position: int): """Log when user had to wait in queue""" self.users_waited_in_queue += 1 self.waiting_times.append(wait_time) avg_wait = sum(self.waiting_times) / len(self.waiting_times) if self.waiting_times else 0 wait_rate = (self.users_waited_in_queue / self.total_connections) * 100 if self.total_connections > 0 else 0 timestamp = time.time() # Human-readable log self._write_log(f"⏳ QUEUE WAIT: {client_id} waited {wait_time:.1f}s (was #{queue_position})") self._write_log(f" 📊 Wait rate: {wait_rate:.1f}% | Avg wait time: {avg_wait:.1f}s") # Structured data log self._write_json_log(self.queue_metrics_file, { "type": "queue_wait", "timestamp": timestamp, "client_id": client_id, "wait_time": wait_time, "queue_position": queue_position, "wait_rate": wait_rate, "avg_wait_time": avg_wait, "users_waited": self.users_waited_in_queue }) def log_gpu_status(self, total_gpus: int, active_gpus: int, available_gpus: int): """Log GPU utilization""" utilization = (active_gpus / total_gpus) * 100 if total_gpus > 0 else 0 self.gpu_utilization_samples.append(utilization) avg_utilization = sum(self.gpu_utilization_samples) / len(self.gpu_utilization_samples) if self.gpu_utilization_samples else 0 timestamp = time.time() # Human-readable log self._write_log(f"🖥️ GPU STATUS: {active_gpus}/{total_gpus} in use ({utilization:.1f}% utilization)") self._write_log(f" 📊 Available: {available_gpus} | Avg utilization: {avg_utilization:.1f}%") # Structured data log self._write_json_log(self.gpu_metrics_file, { "type": "gpu_status", "timestamp": timestamp, "total_gpus": total_gpus, "active_gpus": active_gpus, "available_gpus": available_gpus, "utilization_percent": utilization, "avg_utilization_percent": avg_utilization }) def log_worker_registered(self, worker_id: str, worker_address: str, endpoint: str): """Log when a worker registers""" self._write_log(f"⚙️ WORKER REGISTERED: {worker_id} ({worker_address}) at {endpoint}") def log_worker_disconnected(self, worker_id: str, worker_address: str): """Log when a worker disconnects""" self._write_log(f"⚙️ WORKER DISCONNECTED: {worker_id} ({worker_address})") def log_no_workers_available(self, queue_size: int): """Log critical situation when no workers are available""" self._write_log(f"⚠️ CRITICAL: No GPU workers available! {queue_size} users waiting") self._write_log(" Please check worker processes and GPU availability") def log_queue_limits_applied(self, affected_sessions: int, queue_size: int): """Log when time limits are applied to existing sessions due to queue formation""" timestamp = time.time() # Human-readable log self._write_log(f"🕐 QUEUE LIMITS APPLIED: {affected_sessions} existing sessions now have 60s limits") self._write_log(f" 📊 Reason: Queue formed with {queue_size} waiting users") # Structured data log self._write_json_log(self.queue_metrics_file, { "type": "queue_limits_applied", "timestamp": timestamp, "affected_sessions": affected_sessions, "queue_size": queue_size, "time_limit_applied": 60.0 }) def log_queue_status(self, queue_size: int, maximum_wait: float): """Log queue status""" self.queue_size_samples.append(queue_size) avg_queue_size = sum(self.queue_size_samples) / len(self.queue_size_samples) if self.queue_size_samples else 0 timestamp = time.time() # Always log to structured data for analysis self._write_json_log(self.queue_metrics_file, { "type": "queue_status", "timestamp": timestamp, "queue_size": queue_size, "maximum_wait": maximum_wait, "avg_queue_size": avg_queue_size }) # Only log to human-readable if there's a queue if queue_size > 0: self._write_log(f"📝 QUEUE STATUS: {queue_size} users waiting | Max wait: {maximum_wait:.1f}s") self._write_log(f" 📊 Avg queue size: {avg_queue_size:.1f}") def log_periodic_summary(self): """Log periodic system summary""" uptime = time.time() - self.start_time uptime_hours = uptime / 3600 unique_ips = len(self.ip_addresses) avg_duration = sum(self.session_durations) / len(self.session_durations) if self.session_durations else 0 avg_wait = sum(self.waiting_times) / len(self.waiting_times) if self.waiting_times else 0 avg_utilization = sum(self.gpu_utilization_samples) / len(self.gpu_utilization_samples) if self.gpu_utilization_samples else 0 avg_queue_size = sum(self.queue_size_samples) / len(self.queue_size_samples) if self.queue_size_samples else 0 bypass_rate = (self.users_bypassed_queue / self.total_connections) * 100 if self.total_connections > 0 else 0 self._write_log("") self._write_log("="*60) self._write_log("📊 SYSTEM SUMMARY") self._write_log("="*60) self._write_log(f"⏱️ Uptime: {uptime_hours:.1f} hours") self._write_log(f"🔗 Connections: {self.total_connections} total | {self.active_connections} active | {unique_ips} unique IPs") self._write_log(f"💬 Total interactions: {self.total_interactions}") self._write_log(f"⚡ Queue bypass rate: {bypass_rate:.1f}% ({self.users_bypassed_queue}/{self.total_connections})") self._write_log(f"⏳ Avg waiting time: {avg_wait:.1f}s") self._write_log(f"📝 Avg queue size: {avg_queue_size:.1f}") self._write_log(f"🖥️ Avg GPU utilization: {avg_utilization:.1f}%") self._write_log(f"⏱️ Avg session duration: {avg_duration:.1f}s") self._write_log("") self._write_log("🌍 TOP IP ADDRESSES:") for ip, count in sorted(self.ip_addresses.items(), key=lambda x: x[1], reverse=True)[:10]: self._write_log(f" {ip}: {count} connections") self._write_log("="*60) self._write_log("") # Initialize analytics analytics = SystemAnalytics() class SessionStatus(Enum): QUEUED = "queued" ACTIVE = "active" COMPLETED = "completed" TIMEOUT = "timeout" @dataclass class UserSession: session_id: str client_id: str websocket: WebSocket created_at: float status: SessionStatus worker_id: Optional[str] = None last_activity: Optional[float] = None max_session_time: Optional[float] = None session_limit_start_time: Optional[float] = None # When the time limit was first applied user_has_interacted: bool = False ip_address: Optional[str] = None interaction_count: int = 0 queue_start_time: Optional[float] = None idle_warning_sent: bool = False session_warning_sent: bool = False @dataclass class WorkerInfo: worker_id: str worker_address: str # e.g., "localhost:8001", "192.168.1.100:8002" endpoint: str is_available: bool current_session: Optional[str] = None last_ping: float = 0 class SessionManager: def __init__(self): self.sessions: Dict[str, UserSession] = {} self.workers: Dict[str, WorkerInfo] = {} self.session_queue: List[str] = [] self.active_sessions: Dict[str, str] = {} # session_id -> worker_id self._queue_lock = asyncio.Lock() # Prevent race conditions in queue processing # Configuration self.IDLE_TIMEOUT = 20.0 # When no queue self.QUEUE_WARNING_TIME = 10.0 self.MAX_SESSION_TIME_WITH_QUEUE = 60.0 # When there's a queue self.QUEUE_SESSION_WARNING_TIME = 45.0 # 15 seconds before timeout self.GRACE_PERIOD = 10.0 async def register_worker(self, worker_id: str, worker_address: str, endpoint: str): """Register a new worker""" # Check for duplicate registrations if worker_id in self.workers: logger.warning(f"Worker {worker_id} already registered! Overwriting previous registration.") logger.warning(f"Previous: {self.workers[worker_id].worker_address}, {self.workers[worker_id].endpoint}") logger.warning(f"New: {worker_address}, {endpoint}") self.workers[worker_id] = WorkerInfo( worker_id=worker_id, worker_address=worker_address, endpoint=endpoint, is_available=True, last_ping=time.time() ) logger.info(f"Registered worker {worker_id} ({worker_address}) at {endpoint}") logger.info(f"Total workers now: {len(self.workers)} - {[w.worker_id for w in self.workers.values()]}") # Log worker registration analytics.log_worker_registered(worker_id, worker_address, endpoint) # Log GPU status total_gpus = len(self.workers) active_gpus = len([w for w in self.workers.values() if not w.is_available]) available_gpus = total_gpus - active_gpus analytics.log_gpu_status(total_gpus, active_gpus, available_gpus) async def get_available_worker(self) -> Optional[WorkerInfo]: """Get an available worker""" for worker in self.workers.values(): if worker.is_available and time.time() - worker.last_ping < 20: # Worker ping timeout return worker return None async def add_session_to_queue(self, session: UserSession): """Add a session to the queue""" self.sessions[session.session_id] = session self.session_queue.append(session.session_id) session.status = SessionStatus.QUEUED session.queue_start_time = time.time() logger.info(f"Added session {session.session_id} to queue. Queue size: {len(self.session_queue)}") # Don't apply time limits here - wait until after processing queue # to see if users actually have to wait async def apply_queue_limits_to_existing_sessions(self): """Apply 60-second time limits to existing unlimited sessions when queue forms""" current_time = time.time() affected_sessions = 0 for session_id in list(self.active_sessions.keys()): session = self.sessions.get(session_id) if session and session.max_session_time is None: # Currently unlimited # Give them 60 seconds from now session.max_session_time = 60.0 session.session_limit_start_time = current_time # Track when limit started # Don't reset last_activity - that would break idle timeout tracking session.session_warning_sent = False # Reset warning flag affected_sessions += 1 # Notify the user about the new time limit try: queue_size = len(self.session_queue) await session.websocket.send_json({ "type": "queue_limit_applied", "time_remaining": 60.0, "queue_size": queue_size }) logger.info(f"Applied 60s time limit to existing session {session_id} due to queue formation ({queue_size} users waiting)") except Exception as e: logger.error(f"Failed to notify session {session_id} about queue limit: {e}") if affected_sessions > 0: analytics.log_queue_limits_applied(affected_sessions, len(self.session_queue)) async def remove_time_limits_if_queue_empty(self): """Remove time limits from active sessions when queue becomes empty""" if len(self.session_queue) > 0: return # Queue not empty, don't remove limits removed_limits = 0 for session_id in list(self.active_sessions.keys()): session = self.sessions.get(session_id) if session and session.max_session_time is not None: # Remove time limit session.max_session_time = None session.session_limit_start_time = None session.session_warning_sent = False removed_limits += 1 # Notify user that time limit was removed try: await session.websocket.send_json({ "type": "time_limit_removed", "reason": "queue_empty" }) logger.info(f"Time limit removed from active session {session_id} (queue empty)") except Exception as e: logger.error(f"Failed to notify session {session_id} about time limit removal: {e}") if removed_limits > 0: logger.info(f"Removed time limits from {removed_limits} active sessions (queue became empty)") async def process_queue(self): """Process the session queue""" async with self._queue_lock: # Prevent race conditions # Track if we had any existing active sessions before processing had_active_sessions = len(self.active_sessions) > 0 # Add detailed logging for debugging logger.info(f"Processing queue: {len(self.session_queue)} waiting, {len(self.active_sessions)} active") logger.info(f"Available workers: {[f'{w.worker_id}({w.worker_address})' for w in self.workers.values() if w.is_available]}") logger.info(f"Busy workers: {[f'{w.worker_id}({w.worker_address})' for w in self.workers.values() if not w.is_available]}") while self.session_queue: session_id = self.session_queue[0] session = self.sessions.get(session_id) if not session or session.status != SessionStatus.QUEUED: self.session_queue.pop(0) continue worker = await self.get_available_worker() if not worker: # Log critical situation if no workers are available if len(self.workers) == 0: analytics.log_no_workers_available(len(self.session_queue)) logger.info(f"No available workers for session {session_id}. Queue processing stopped.") break # No available workers # Calculate wait time wait_time = time.time() - session.queue_start_time if session.queue_start_time else 0 queue_position = self.session_queue.index(session_id) + 1 # Assign session to worker self.session_queue.pop(0) session.status = SessionStatus.ACTIVE session.worker_id = worker.worker_id session.last_activity = time.time() # Set session time limit based on queue status AFTER processing if len(self.session_queue) > 0: session.max_session_time = self.MAX_SESSION_TIME_WITH_QUEUE session.session_limit_start_time = time.time() # Track when limit started worker.is_available = False worker.current_session = session_id self.active_sessions[session_id] = worker.worker_id logger.info(f"Assigned session {session_id} to worker {worker.worker_id}") logger.info(f"Active sessions now: {len(self.active_sessions)}, Available workers: {len([w for w in self.workers.values() if w.is_available])}") # Log analytics if wait_time > 0: analytics.log_queue_wait(session.client_id, wait_time, queue_position) else: analytics.log_queue_bypass(session.client_id) # Log GPU status total_gpus = len(self.workers) active_gpus = len([w for w in self.workers.values() if not w.is_available]) available_gpus = total_gpus - active_gpus analytics.log_gpu_status(total_gpus, active_gpus, available_gpus) # Initialize session on worker with client_id for logging try: async with aiohttp.ClientSession() as client_session: await client_session.post(f"{worker.endpoint}/init_session", json={ "session_id": session_id, "client_id": session.client_id }) except Exception as e: logger.error(f"Failed to initialize session on worker {worker.worker_id}: {e}") # Notify user that their session is starting await self.notify_session_start(session) # Start session monitoring asyncio.create_task(self.monitor_active_session(session_id)) # After processing queue, if there are still users waiting AND we had existing active sessions, # apply time limits to those existing sessions if len(self.session_queue) > 0 and had_active_sessions: await self.apply_queue_limits_to_existing_sessions() # If queue became empty and there are active sessions with time limits, remove them elif len(self.session_queue) == 0: await self.remove_time_limits_if_queue_empty() async def notify_session_start(self, session: UserSession): """Notify user that their session is starting""" try: await session.websocket.send_json({ "type": "session_start", "worker_id": session.worker_id, "max_session_time": session.max_session_time }) except Exception as e: logger.error(f"Failed to notify session start for {session.session_id}: {e}") async def monitor_active_session(self, session_id: str): """Monitor an active session for timeouts""" session = self.sessions.get(session_id) if not session: return try: while session.status == SessionStatus.ACTIVE: current_time = time.time() # Check timeouts - both session limit AND idle timeout can apply session_timeout = False idle_timeout = False # Check session time limit (when queue exists) if session.max_session_time: # Use session_limit_start_time for absolute timeout, fall back to last_activity if not set start_time = session.session_limit_start_time if session.session_limit_start_time else session.last_activity elapsed = current_time - start_time if start_time else 0 remaining = session.max_session_time - elapsed # Send warning at 15 seconds before timeout (only once) if remaining <= 15 and remaining > 10 and not session.session_warning_sent: await session.websocket.send_json({ "type": "session_warning", "time_remaining": remaining, "queue_size": len(self.session_queue) }) session.session_warning_sent = True logger.info(f"Session warning sent to {session_id}, time remaining: {remaining:.1f}s") # Grace period handling elif remaining <= 10 and remaining > 0: # Check if queue is empty - if so, extend session if len(self.session_queue) == 0: session.max_session_time = None # Remove time limit session.session_limit_start_time = None # Clear time limit start time session.session_warning_sent = False # Reset warning since limit removed await session.websocket.send_json({ "type": "time_limit_removed", "reason": "queue_empty" }) logger.info(f"Session time limit removed for {session_id} (queue empty)") else: await session.websocket.send_json({ "type": "grace_period", "time_remaining": remaining, "queue_size": len(self.session_queue) }) # Session timeout elif remaining <= 0: session_timeout = True # Check idle timeout (always check when user has interacted) if session.last_activity: idle_time = current_time - session.last_activity if idle_time >= self.IDLE_TIMEOUT: idle_timeout = True elif idle_time >= self.QUEUE_WARNING_TIME and not session.idle_warning_sent: # Send idle warning (even when session limit exists) await session.websocket.send_json({ "type": "idle_warning", "time_remaining": self.IDLE_TIMEOUT - idle_time }) session.idle_warning_sent = True reason = "with session limit" if session.max_session_time else "no session limit" logger.info(f"Idle warning sent to {session_id}, time remaining: {self.IDLE_TIMEOUT - idle_time:.1f}s ({reason})") # End session if either timeout triggered if session_timeout or idle_timeout: timeout_reason = "session_limit" if session_timeout else "idle" logger.info(f"Ending session {session_id} due to {timeout_reason} timeout") await self.end_session(session_id, SessionStatus.TIMEOUT) return await asyncio.sleep(1) # Check every second except Exception as e: logger.error(f"Error monitoring session {session_id}: {e}") await self.end_session(session_id, SessionStatus.COMPLETED) async def end_session(self, session_id: str, status: SessionStatus): """End a session and free up the worker""" session = self.sessions.get(session_id) if not session: return session.status = status # Calculate session duration duration = time.time() - session.created_at # Log analytics reason = "timeout" if status == SessionStatus.TIMEOUT else "normal" analytics.log_connection_closed(session.client_id, duration, session.interaction_count, reason) # Free up the worker if session.worker_id and session.worker_id in self.workers: worker = self.workers[session.worker_id] if not worker.is_available: # Only log if worker was actually busy logger.info(f"Freeing worker {worker.worker_id} from session {session_id}") else: logger.warning(f"Worker {worker.worker_id} was already available when freeing from session {session_id}") worker.is_available = True worker.current_session = None # Log GPU status total_gpus = len(self.workers) active_gpus = len([w for w in self.workers.values() if not w.is_available]) available_gpus = total_gpus - active_gpus analytics.log_gpu_status(total_gpus, active_gpus, available_gpus) # Notify worker to clean up try: async with aiohttp.ClientSession() as client_session: await client_session.post(f"{worker.endpoint}/end_session", json={"session_id": session_id}) except Exception as e: logger.error(f"Failed to notify worker {worker.worker_id} of session end: {e}") # Remove from active sessions if session_id in self.active_sessions: del self.active_sessions[session_id] logger.info(f"Ended session {session_id} with status {status}") # Validate system state consistency await self._validate_system_state() # Process next in queue asyncio.create_task(self.process_queue()) async def update_queue_info(self): """Send queue information to waiting users""" for i, session_id in enumerate(self.session_queue): session = self.sessions.get(session_id) if session and session.status == SessionStatus.QUEUED: try: # Calculate maximum possible wait time maximum_wait = self._calculate_maximum_wait_time(i + 1) await session.websocket.send_json({ "type": "queue_update", "position": i + 1, "total_waiting": len(self.session_queue), "maximum_wait_seconds": maximum_wait, "active_sessions": len(self.active_sessions), "available_workers": len([w for w in self.workers.values() if w.is_available]) }) except Exception as e: logger.error(f"Failed to send queue update to session {session_id}: {e}") # Log queue status if there's a queue if self.session_queue: maximum_wait = self._calculate_maximum_wait_time(1) analytics.log_queue_status(len(self.session_queue), maximum_wait) def _calculate_maximum_wait_time(self, position_in_queue: int) -> float: """Calculate realistic wait time based on actual session remaining times""" available_workers = len([w for w in self.workers.values() if w.is_available]) # If there are available workers, no wait time if available_workers > 0: return 0 # Calculate wait time based on position and worker count num_workers = len(self.workers) if num_workers == 0: return 999 # No workers available # Get actual remaining times for active sessions current_time = time.time() active_session_remaining_times = [] for session_id in self.active_sessions: session = self.sessions.get(session_id) if session: # Only consider session time limit (hard limit), not idle timeout # Users can reset idle timeout by moving mouse, but session limit is fixed if session.max_session_time and session.session_limit_start_time: elapsed = current_time - session.session_limit_start_time remaining_time = max(0, session.max_session_time - elapsed) else: # No session limit set (shouldn't happen when queue exists, but fallback) remaining_time = self.MAX_SESSION_TIME_WITH_QUEUE active_session_remaining_times.append(remaining_time) # Sort remaining times (shortest first) active_session_remaining_times.sort() # Calculate when this user will get a worker if position_in_queue <= len(active_session_remaining_times): # User will get a worker when the Nth session ends estimated_wait = active_session_remaining_times[position_in_queue - 1] logger.info(f"Queue position {position_in_queue}: Will get worker when session #{position_in_queue} ends in {estimated_wait:.1f}s") logger.info(f"Active session remaining times: {[f'{t:.1f}s' for t in active_session_remaining_times]}") else: # More people in queue than active sessions # Need to wait for multiple "waves" to complete full_waves = (position_in_queue - 1) // num_workers position_in_final_wave = (position_in_queue - 1) % num_workers # Wait for current sessions to end, then full waves, then partial wave if active_session_remaining_times: first_wave_time = max(active_session_remaining_times) else: first_wave_time = self.MAX_SESSION_TIME_WITH_QUEUE additional_wave_time = full_waves * self.MAX_SESSION_TIME_WITH_QUEUE estimated_wait = first_wave_time + additional_wave_time logger.info(f"Queue position {position_in_queue}: Need {full_waves} full waves + current sessions. Wait: {estimated_wait:.1f}s") return estimated_wait async def handle_user_activity(self, session_id: str): """Update user activity timestamp and reset warning flags""" session = self.sessions.get(session_id) if session: old_time = session.last_activity session.interaction_count += 1 # Always update last_activity for idle detection # Session time limits are now tracked separately via session_limit_start_time session.last_activity = time.time() if session.max_session_time is not None: logger.info(f"Activity detected for session {session_id} - idle timer reset (session limit still active)") else: logger.info(f"Activity detected for session {session_id} - idle timer reset (no session limit)") # Reset warning flags if user activity detected after warnings warning_reset = False if session.idle_warning_sent: logger.info(f"User activity detected, resetting idle warning for session {session_id}") session.idle_warning_sent = False warning_reset = True if session.session_warning_sent: logger.info(f"User activity detected, resetting session warning for session {session_id}") session.session_warning_sent = False warning_reset = True # Notify client that activity was detected and warnings are reset if warning_reset: try: # Include current session time remaining if session limit is active message = {"type": "activity_reset"} if session.max_session_time and session.session_limit_start_time: elapsed = time.time() - session.session_limit_start_time remaining = max(0, session.max_session_time - elapsed) if remaining > 0: message["session_time_remaining"] = remaining message["queue_size"] = len(self.session_queue) await session.websocket.send_json(message) logger.info(f"Activity reset message sent to session {session_id}") except Exception as e: logger.error(f"Failed to send activity reset to session {session_id}: {e}") if not session.user_has_interacted: session.user_has_interacted = True logger.info(f"User started interacting in session {session_id}") async def _validate_system_state(self): """Validate system state consistency for debugging""" try: # Count active sessions active_sessions_count = len(self.active_sessions) busy_workers_count = len([w for w in self.workers.values() if not w.is_available]) # Check for inconsistencies if active_sessions_count != busy_workers_count: logger.error(f"INCONSISTENCY: Active sessions ({active_sessions_count}) != Busy workers ({busy_workers_count})") logger.error(f"Active sessions: {list(self.active_sessions.keys())}") logger.error(f"Busy workers: {[w.worker_id for w in self.workers.values() if not w.is_available]}") # Log detailed state for session_id, worker_id in self.active_sessions.items(): session = self.sessions.get(session_id) worker = self.workers.get(worker_id) logger.error(f"Session {session_id}: status={session.status if session else 'MISSING'}, worker={worker_id}") if worker: logger.error(f"Worker {worker_id}: available={worker.is_available}, current_session={worker.current_session}") # Check for orphaned workers for worker in self.workers.values(): if not worker.is_available and worker.current_session not in self.active_sessions: logger.error(f"ORPHANED WORKER: {worker.worker_id} is busy but session {worker.current_session} not in active_sessions") # Check for sessions without workers for session_id in self.active_sessions: session = self.sessions.get(session_id) if session and session.status == SessionStatus.ACTIVE: worker = self.workers.get(session.worker_id) if not worker or worker.is_available: logger.error(f"ACTIVE SESSION WITHOUT WORKER: {session_id} has no busy worker assigned") except Exception as e: logger.error(f"Error in system state validation: {e}") async def _handle_worker_failure(self, failed_worker_id: str): """Handle sessions when a worker fails - end sessions and put users back in queue""" logger.warning(f"Handling failure of worker {failed_worker_id}") # Find all sessions assigned to this worker failed_sessions = [] for session_id, worker_id in list(self.active_sessions.items()): if worker_id == failed_worker_id: failed_sessions.append(session_id) logger.warning(f"Found {len(failed_sessions)} sessions on failed worker {failed_worker_id}") for session_id in failed_sessions: session = self.sessions.get(session_id) if session: logger.info(f"Recovering session {session_id} from failed worker") # Notify user about the worker failure try: await session.websocket.send_json({ "type": "worker_failure", "message": "GPU worker failed. Reconnecting you to a healthy worker..." }) except Exception as e: logger.error(f"Failed to notify session {session_id} about worker failure: {e}") # Remove from active sessions if session_id in self.active_sessions: del self.active_sessions[session_id] # Reset session state and put back in queue session.status = SessionStatus.QUEUED session.worker_id = None session.queue_start_time = time.time() session.max_session_time = None # Reset time limits session.session_limit_start_time = None session.session_warning_sent = False session.idle_warning_sent = False # Add back to front of queue (they were already active) if session_id not in self.session_queue: self.session_queue.insert(0, session_id) logger.info(f"Added session {session_id} to front of queue for recovery") # Process queue to reassign recovered sessions to healthy workers if failed_sessions: logger.info(f"Processing queue to reassign {len(failed_sessions)} recovered sessions") await self.process_queue() async def _forward_to_worker(self, worker: WorkerInfo, session_id: str, data: dict): """Forward input to worker asynchronously""" try: async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as client_session: async with client_session.post( f"{worker.endpoint}/process_input", json={ "session_id": session_id, "data": data } ) as response: if response.status != 200: logger.error(f"Worker {worker.worker_id} returned status {response.status}") except asyncio.TimeoutError: logger.error(f"Worker {worker.worker_id} timeout - may be unresponsive") # Mark worker as potentially dead for faster detection worker.last_ping = 0 # This will cause it to be removed on next health check except Exception as e: logger.error(f"Error forwarding to worker {worker.worker_id}: {e}") # Mark worker as potentially dead for faster detection worker.last_ping = 0 # Global session manager session_manager = SessionManager() app = FastAPI() app.mount("/static", StaticFiles(directory="static"), name="static") # Global connection counter like in main.py connection_counter = 0 @app.get("/") async def get(): return HTMLResponse(open("static/index.html").read()) @app.post("/register_worker") async def register_worker(worker_info: dict): """Endpoint for workers to register themselves""" logger.info(f"📥 Received worker registration request") logger.info(f"📊 Worker info: {worker_info}") try: await session_manager.register_worker( worker_info["worker_id"], worker_info["worker_address"], worker_info["endpoint"] ) logger.info(f"✅ Successfully processed worker registration") return {"status": "registered"} except Exception as e: logger.error(f"❌ Failed to register worker: {e}") import traceback logger.error(f"🔍 Full traceback: {traceback.format_exc()}") raise @app.post("/worker_ping") async def worker_ping(worker_info: dict): """Endpoint for workers to ping their availability""" worker_id = worker_info["worker_id"] if worker_id in session_manager.workers: session_manager.workers[worker_id].last_ping = time.time() session_manager.workers[worker_id].is_available = worker_info.get("is_available", True) return {"status": "ok"} @app.post("/worker_result") async def worker_result(result_data: dict): """Endpoint for workers to send back processing results""" session_id = result_data.get("session_id") worker_id = result_data.get("worker_id") result = result_data.get("result") if not session_id or not result: raise HTTPException(status_code=400, detail="Missing session_id or result") # Find the session and send result to the WebSocket session = session_manager.sessions.get(session_id) if session and session.status == SessionStatus.ACTIVE: try: await session.websocket.send_json(result) logger.info(f"Sent result to session {session_id}") except Exception as e: logger.error(f"Failed to send result to session {session_id}: {e}") else: logger.warning(f"Could not find active session {session_id} for result") return {"status": "ok"} @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): global connection_counter await websocket.accept() # Extract client IP address client_ip = "unknown" if websocket.client and hasattr(websocket.client, 'host'): client_ip = websocket.client.host elif hasattr(websocket, 'headers'): # Try to get real IP from headers (for proxy setups) client_ip = websocket.headers.get('x-forwarded-for', websocket.headers.get('x-real-ip', websocket.headers.get('cf-connecting-ip', 'unknown'))) if ',' in client_ip: client_ip = client_ip.split(',')[0].strip() # Create session with connection counter like in main.py connection_counter += 1 session_id = str(uuid.uuid4()) client_id = f"{int(time.time())}_{connection_counter}" session = UserSession( session_id=session_id, client_id=client_id, websocket=websocket, created_at=time.time(), status=SessionStatus.QUEUED, ip_address=client_ip ) logger.info(f"New WebSocket connection: {client_id} from {client_ip}") # Log new connection analytics analytics.log_new_connection(client_id, client_ip) try: # Add to queue await session_manager.add_session_to_queue(session) # Try to process queue immediately await session_manager.process_queue() # Send initial queue status if session.status == SessionStatus.QUEUED: await session_manager.update_queue_info() # Main message loop while True: try: data = await websocket.receive_json() # Update activity only for real user inputs, not auto inputs if not data.get("is_auto_input", False): # Log stay-connected attempts for monitoring if data.get("is_stay_connected", False): logger.info(f"Stay-connected ping from session {session_id} - idle timer will be reset") await session_manager.handle_user_activity(session_id) # Handle different message types if data.get("type") == "heartbeat": await websocket.send_json({"type": "heartbeat_response"}) continue # If session is active, forward to worker if session.status == SessionStatus.ACTIVE and session.worker_id: worker = session_manager.workers.get(session.worker_id) if worker: try: # Forward message to worker (don't wait for response for regular inputs) # The worker will send results back asynchronously via /worker_result asyncio.create_task(session_manager._forward_to_worker(worker, session_id, data)) except Exception as e: logger.error(f"Error forwarding to worker: {e}") # Handle control messages (these need synchronous responses) elif data.get("type") in ["reset", "update_sampling_steps", "update_use_rnn", "get_settings"]: if session.status == SessionStatus.ACTIVE and session.worker_id: worker = session_manager.workers.get(session.worker_id) if worker: try: async with aiohttp.ClientSession() as client_session: async with client_session.post( f"{worker.endpoint}/process_input", json={ "session_id": session_id, "data": data } ) as response: if response.status == 200: result = await response.json() await websocket.send_json(result) else: logger.error(f"Worker returned status {response.status}") except Exception as e: logger.error(f"Error forwarding control message: {e}") else: # Send appropriate response for queued users await websocket.send_json({ "type": "error", "message": "Session not active yet. Please wait in queue." }) except asyncio.TimeoutError: logger.info("WebSocket connection timed out") break except WebSocketDisconnect: logger.info(f"WebSocket disconnected: {client_id}") break except Exception as e: logger.error(f"Error in WebSocket connection {client_id}: {e}") import traceback traceback.print_exc() finally: # Clean up session if session_id in session_manager.sessions: await session_manager.end_session(session_id, SessionStatus.COMPLETED) del session_manager.sessions[session_id] logger.info(f"WebSocket connection closed: {client_id}") # Background task to periodically update queue info async def periodic_queue_update(): while True: try: await session_manager.update_queue_info() await asyncio.sleep(2) # Update every 2 seconds for responsive experience except Exception as e: logger.error(f"Error in periodic queue update: {e}") # Background task to periodically validate system state async def periodic_system_validation(): while True: try: await asyncio.sleep(10) # Validate every 10 seconds await session_manager._validate_system_state() except Exception as e: logger.error(f"Error in periodic system validation: {e}") # Background task to periodically log analytics summary async def periodic_analytics_summary(): while True: try: await asyncio.sleep(300) # Log summary every 5 minutes analytics.log_periodic_summary() except Exception as e: logger.error(f"Error in periodic analytics summary: {e}") # Background task to check worker health async def periodic_worker_health_check(): while True: try: await asyncio.sleep(15) # Check every 15 seconds current_time = time.time() disconnected_workers = [] for worker_id, worker in list(session_manager.workers.items()): if current_time - worker.last_ping > 20: # 20 second timeout disconnected_workers.append((worker_id, worker.worker_address)) for worker_id, worker_address in disconnected_workers: analytics.log_worker_disconnected(worker_id, worker_address) # Handle any active sessions on this dead worker await session_manager._handle_worker_failure(worker_id) del session_manager.workers[worker_id] logger.warning(f"Removed disconnected worker {worker_id} ({worker_address})") if disconnected_workers: # Log updated GPU status total_gpus = len(session_manager.workers) active_gpus = len([w for w in session_manager.workers.values() if not w.is_available]) available_gpus = total_gpus - active_gpus analytics.log_gpu_status(total_gpus, active_gpus, available_gpus) except Exception as e: logger.error(f"Error in periodic worker health check: {e}") @app.on_event("startup") async def startup_event(): logger.info("🚀 Dispatcher startup event triggered") # Start background tasks asyncio.create_task(periodic_queue_update()) asyncio.create_task(periodic_system_validation()) asyncio.create_task(periodic_analytics_summary()) asyncio.create_task(periodic_worker_health_check()) # Log initial system status analytics._write_log("🚀 System initialized and ready to accept connections") analytics._write_log(" Waiting for GPU workers to register...") logger.info("✅ Dispatcher startup complete - ready to accept worker registrations") @app.on_event("shutdown") async def shutdown_event(): # Log final system summary analytics._write_log("") analytics._write_log("🛑 System shutting down...") analytics.log_periodic_summary() analytics._write_log("System shutdown complete.") if __name__ == "__main__": import uvicorn import argparse parser = argparse.ArgumentParser(description="Dispatcher for Neural OS") parser.add_argument("--port", type=int, default=7860, help="Port to run the dispatcher on") args = parser.parse_args() logger.info(f"🌐 Starting dispatcher on 0.0.0.0:{args.port}") logger.info(f"🔗 Dispatcher will be available at http://localhost:{args.port}") try: uvicorn.run(app, host="0.0.0.0", port=args.port) except Exception as e: logger.error(f"❌ Failed to start dispatcher: {e}") import traceback logger.error(f"🔍 Full traceback: {traceback.format_exc()}") raise