Spaces:
Runtime error
Runtime error
File size: 9,426 Bytes
602eeec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 | """
Simplified Session Manager (inspired by Udbhav's approach)
Features:
- In-memory session tracking with thread-safe RLock
- Simple filesystem-based storage (one directory per session)
- Persistent session index (JSON) and alias map
- Automatic cleanup of idle sessions (>6 hours)
- Human-readable aliases for sessions
"""
import os
import json
import shutil
import threading
import uuid
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, Optional
# Base directory for all sessions (HF persistent storage)
BASE_DATA_DIR = Path("/tmp/outputs/sessions")
# Session index file (persisted to disk)
SESSIONS_INDEX_FILE = BASE_DATA_DIR / "sessions_index.json"
# Alias map file (maps friendly names to session IDs)
ALIAS_MAP_FILE = BASE_DATA_DIR / "alias_map.json"
class SimpleSessionManager:
"""
Thread-safe session manager with persistence and auto-cleanup.
"""
def __init__(self):
self.sessions: Dict[str, dict] = {}
self.alias_map: Dict[str, str] = {} # alias -> session_id
self._lock = threading.RLock()
self._cleanup_thread = None
# Ensure directories exist
BASE_DATA_DIR.mkdir(parents=True, exist_ok=True)
print(f"[SESSION] Base directory: {BASE_DATA_DIR}")
# Load persisted data
self._load_index()
self._load_alias_map()
# Start cleanup thread
self._start_cleanup_thread()
def _load_index(self):
"""Load sessions index from disk."""
if SESSIONS_INDEX_FILE.exists():
try:
with open(SESSIONS_INDEX_FILE, 'r') as f:
data = json.load(f)
self.sessions = data
print(f"[SESSION] Loaded {len(self.sessions)} sessions from index")
except Exception as e:
print(f"[SESSION] Error loading index: {e}")
def _load_alias_map(self):
"""Load alias map from disk."""
if ALIAS_MAP_FILE.exists():
try:
with open(ALIAS_MAP_FILE, 'r') as f:
self.alias_map = json.load(f)
print(f"[SESSION] Loaded {len(self.alias_map)} aliases")
except Exception as e:
print(f"[SESSION] Error loading alias map: {e}")
def _save_index(self):
"""Save sessions index to disk."""
try:
SESSIONS_INDEX_FILE.parent.mkdir(parents=True, exist_ok=True)
with open(SESSIONS_INDEX_FILE, 'w') as f:
json.dump(self.sessions, f, indent=2, default=str)
except Exception as e:
print(f"[SESSION] Error saving index: {e}")
def _save_alias_map(self):
"""Save alias map to disk."""
try:
ALIAS_MAP_FILE.parent.mkdir(parents=True, exist_ok=True)
with open(ALIAS_MAP_FILE, 'w') as f:
json.dump(self.alias_map, f, indent=2)
except Exception as e:
print(f"[SESSION] Error saving alias map: {e}")
def _new_session(self, session_id: Optional[str] = None, alias: Optional[str] = None) -> Dict:
"""Create a new session directory structure."""
if session_id is None:
session_id = str(uuid.uuid4())
session_root = BASE_DATA_DIR / f"user_{session_id}"
# Create subdirectories
subdirs = {
"data": session_root / "data",
"cache": session_root / "cache",
"results": session_root / "results",
}
for dir_path in subdirs.values():
dir_path.mkdir(parents=True, exist_ok=True)
session_info = {
"session_id": session_id,
"root": str(session_root),
"subdirs": {k: str(v) for k, v in subdirs.items()},
"created_at": datetime.now().isoformat(),
"last_accessed": datetime.now().isoformat(),
"alias": alias,
}
print(f"[SESSION] Created new session: {session_id} (alias: {alias})")
return session_info
def create_session(self, alias: Optional[str] = None) -> Dict:
"""
Create a new session with optional alias.
Args:
alias: Human-readable name for the session
Returns:
Session info dict
"""
with self._lock:
# If alias exists, return existing session
if alias and alias in self.alias_map:
session_id = self.alias_map[alias]
if session_id in self.sessions:
print(f"[SESSION] Session '{alias}' already exists: {session_id}")
self.sessions[session_id]["last_accessed"] = datetime.now().isoformat()
self._save_index()
return self.sessions[session_id]
# Create new session
session_id = str(uuid.uuid4())
session_info = self._new_session(session_id, alias)
self.sessions[session_id] = session_info
if alias:
self.alias_map[alias] = session_id
self._save_alias_map()
self._save_index()
return session_info
def get_session(self, session_id: Optional[str] = None, alias: Optional[str] = None) -> Optional[Dict]:
"""
Get session by ID or alias.
Args:
session_id: Session UUID
alias: Human-readable session name
Returns:
Session info dict or None
"""
with self._lock:
# Lookup by alias first
if alias:
session_id = self.alias_map.get(alias)
if not session_id:
return None
session_info = self.sessions.get(session_id)
if session_info:
session_info["last_accessed"] = datetime.now().isoformat()
self._save_index()
return session_info
def list_sessions(self) -> Dict[str, Dict]:
"""List all active sessions."""
with self._lock:
return {alias: self.sessions[sid] for alias, sid in self.alias_map.items() if sid in self.sessions}
def delete_session(self, session_id: str) -> bool:
"""Delete a session and its files."""
with self._lock:
session_info = self.sessions.pop(session_id, None)
if not session_info:
return False
# Remove alias mapping
for alias, sid in list(self.alias_map.items()):
if sid == session_id:
self.alias_map.pop(alias)
# Remove directory
session_root = Path(session_info["root"])
try:
if session_root.exists():
shutil.rmtree(session_root)
print(f"[SESSION] Deleted session directory: {session_root}")
except Exception as e:
print(f"[SESSION] Error deleting {session_root}: {e}")
self._save_index()
self._save_alias_map()
return True
def _cleanup_old(self, max_age_hours: int = 6):
"""Remove sessions idle longer than max_age_hours."""
now = datetime.now()
to_remove = []
with self._lock:
for session_id, session_info in self.sessions.items():
last_accessed = datetime.fromisoformat(session_info["last_accessed"])
if now - last_accessed > timedelta(hours=max_age_hours):
to_remove.append(session_id)
for session_id in to_remove:
print(f"[SESSION] Cleaning up idle session: {session_id}")
self.delete_session(session_id)
def _start_cleanup_thread(self):
"""Start background cleanup thread."""
if self._cleanup_thread and self._cleanup_thread.is_alive():
return
def loop():
while True:
try:
# Run cleanup every hour, remove sessions idle > 6 hours
time.sleep(3600)
self._cleanup_old(max_age_hours=6)
except Exception as e:
print(f"[SESSION] Cleanup error: {e}")
self._cleanup_thread = threading.Thread(target=loop, daemon=True)
self._cleanup_thread.start()
print("[SESSION] Cleanup thread started")
# Global session manager instance
SESSION_MANAGER = SimpleSessionManager()
def get_or_create_session(alias: Optional[str] = None) -> Dict:
"""Convenience function to get or create a session."""
if alias:
session = SESSION_MANAGER.get_session(alias=alias)
if session:
return session
return SESSION_MANAGER.create_session(alias=alias)
def get_session(session_id: str) -> Optional[Dict]:
"""Get session by ID."""
return SESSION_MANAGER.get_session(session_id=session_id)
def get_session_by_alias(alias: str) -> Optional[Dict]:
"""Get session by alias."""
return SESSION_MANAGER.get_session(alias=alias)
def list_all_sessions() -> Dict[str, Dict]:
"""List all sessions."""
return SESSION_MANAGER.list_sessions()
|