425 lines
14 KiB
Python
425 lines
14 KiB
Python
"""
|
|
Shared Data Manager for UI Stability Fix
|
|
|
|
Manages data sharing between processes through files with proper locking
|
|
and atomic operations to prevent corruption and conflicts.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
import tempfile
|
|
import platform
|
|
from datetime import datetime
|
|
from dataclasses import dataclass, asdict
|
|
from typing import Dict, Any, Optional, Union
|
|
from pathlib import Path
|
|
import logging
|
|
|
|
# Windows-compatible file locking
|
|
if platform.system() == "Windows":
|
|
import msvcrt
|
|
else:
|
|
import fcntl
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
@dataclass
|
|
class ProcessStatus:
|
|
"""Model for process status information"""
|
|
name: str
|
|
pid: int
|
|
status: str # 'running', 'stopped', 'error'
|
|
start_time: datetime
|
|
last_heartbeat: datetime
|
|
memory_usage: float
|
|
cpu_usage: float
|
|
error_message: Optional[str] = None
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary with datetime serialization"""
|
|
data = asdict(self)
|
|
data['start_time'] = self.start_time.isoformat()
|
|
data['last_heartbeat'] = self.last_heartbeat.isoformat()
|
|
return data
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> 'ProcessStatus':
|
|
"""Create from dictionary with datetime deserialization"""
|
|
data['start_time'] = datetime.fromisoformat(data['start_time'])
|
|
data['last_heartbeat'] = datetime.fromisoformat(data['last_heartbeat'])
|
|
return cls(**data)
|
|
|
|
@dataclass
|
|
class TrainingStatus:
|
|
"""Model for training status information"""
|
|
is_running: bool
|
|
current_epoch: int
|
|
total_epochs: int
|
|
loss: float
|
|
accuracy: float
|
|
last_update: datetime
|
|
model_path: str
|
|
error_message: Optional[str] = None
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary with datetime serialization"""
|
|
data = asdict(self)
|
|
data['last_update'] = self.last_update.isoformat()
|
|
return data
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> 'TrainingStatus':
|
|
"""Create from dictionary with datetime deserialization"""
|
|
data['last_update'] = datetime.fromisoformat(data['last_update'])
|
|
return cls(**data)
|
|
|
|
@dataclass
|
|
class DashboardState:
|
|
"""Model for dashboard state information"""
|
|
is_connected: bool
|
|
last_data_update: datetime
|
|
active_connections: int
|
|
error_count: int
|
|
performance_metrics: Dict[str, float]
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary with datetime serialization"""
|
|
data = asdict(self)
|
|
data['last_data_update'] = self.last_data_update.isoformat()
|
|
return data
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> 'DashboardState':
|
|
"""Create from dictionary with datetime deserialization"""
|
|
data['last_data_update'] = datetime.fromisoformat(data['last_data_update'])
|
|
return cls(**data)
|
|
|
|
|
|
class SharedDataManager:
|
|
"""
|
|
Manages data sharing between processes through files with proper locking
|
|
and atomic operations to prevent corruption and conflicts.
|
|
"""
|
|
|
|
def __init__(self, data_dir: str = "shared_data"):
|
|
"""
|
|
Initialize the shared data manager
|
|
|
|
Args:
|
|
data_dir: Directory to store shared data files
|
|
"""
|
|
self.data_dir = Path(data_dir)
|
|
self.data_dir.mkdir(exist_ok=True)
|
|
|
|
# Define file paths for different data types
|
|
self.training_status_file = self.data_dir / "training_status.json"
|
|
self.dashboard_state_file = self.data_dir / "dashboard_state.json"
|
|
self.process_status_file = self.data_dir / "process_status.json"
|
|
self.market_data_file = self.data_dir / "market_data.json"
|
|
self.model_metrics_file = self.data_dir / "model_metrics.json"
|
|
|
|
logger.info(f"SharedDataManager initialized with data directory: {self.data_dir}")
|
|
|
|
def _lock_file(self, file_handle, exclusive=True):
|
|
"""Cross-platform file locking"""
|
|
if platform.system() == "Windows":
|
|
# Windows file locking
|
|
try:
|
|
if exclusive:
|
|
msvcrt.locking(file_handle.fileno(), msvcrt.LK_LOCK, 1)
|
|
else:
|
|
msvcrt.locking(file_handle.fileno(), msvcrt.LK_LOCK, 1)
|
|
except IOError:
|
|
pass # File locking may not be available in all scenarios
|
|
else:
|
|
# Unix file locking
|
|
lock_type = fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH
|
|
fcntl.flock(file_handle.fileno(), lock_type)
|
|
|
|
def _unlock_file(self, file_handle):
|
|
"""Cross-platform file unlocking"""
|
|
if platform.system() == "Windows":
|
|
try:
|
|
msvcrt.locking(file_handle.fileno(), msvcrt.LK_UNLCK, 1)
|
|
except IOError:
|
|
pass
|
|
else:
|
|
fcntl.flock(file_handle.fileno(), fcntl.LOCK_UN)
|
|
|
|
def _write_json_atomic(self, file_path: Path, data: Dict[str, Any]) -> None:
|
|
"""
|
|
Write JSON data atomically with file locking
|
|
|
|
Args:
|
|
file_path: Path to the file to write
|
|
data: Data to write as JSON
|
|
"""
|
|
temp_path = None
|
|
try:
|
|
# Create temporary file in the same directory
|
|
temp_fd, temp_path = tempfile.mkstemp(
|
|
dir=file_path.parent,
|
|
prefix=f".{file_path.name}.",
|
|
suffix=".tmp"
|
|
)
|
|
|
|
with os.fdopen(temp_fd, 'w') as temp_file:
|
|
# Lock the temporary file
|
|
self._lock_file(temp_file, exclusive=True)
|
|
|
|
# Write data with proper formatting
|
|
json.dump(data, temp_file, indent=2, default=str)
|
|
temp_file.flush()
|
|
os.fsync(temp_file.fileno())
|
|
|
|
# Unlock before closing
|
|
self._unlock_file(temp_file)
|
|
|
|
# Atomically replace the original file
|
|
os.replace(temp_path, file_path)
|
|
logger.debug(f"Successfully wrote data to {file_path}")
|
|
|
|
except Exception as e:
|
|
# Clean up temporary file if it exists
|
|
if temp_path:
|
|
try:
|
|
os.unlink(temp_path)
|
|
except:
|
|
pass
|
|
logger.error(f"Failed to write data to {file_path}: {e}")
|
|
raise
|
|
|
|
def _read_json_safe(self, file_path: Path) -> Dict[str, Any]:
|
|
"""
|
|
Read JSON data safely with file locking
|
|
|
|
Args:
|
|
file_path: Path to the file to read
|
|
|
|
Returns:
|
|
Dictionary containing the JSON data
|
|
"""
|
|
if not file_path.exists():
|
|
logger.debug(f"File {file_path} does not exist, returning empty dict")
|
|
return {}
|
|
|
|
try:
|
|
with open(file_path, 'r') as file:
|
|
# Lock the file for reading
|
|
self._lock_file(file, exclusive=False)
|
|
data = json.load(file)
|
|
self._unlock_file(file)
|
|
logger.debug(f"Successfully read data from {file_path}")
|
|
return data
|
|
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Invalid JSON in {file_path}: {e}")
|
|
return {}
|
|
except Exception as e:
|
|
logger.error(f"Failed to read data from {file_path}: {e}")
|
|
return {}
|
|
|
|
def write_training_status(self, status: TrainingStatus) -> None:
|
|
"""
|
|
Write training status to shared file
|
|
|
|
Args:
|
|
status: TrainingStatus object to write
|
|
"""
|
|
try:
|
|
data = status.to_dict()
|
|
self._write_json_atomic(self.training_status_file, data)
|
|
logger.debug("Training status written successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to write training status: {e}")
|
|
raise
|
|
|
|
def read_training_status(self) -> Optional[TrainingStatus]:
|
|
"""
|
|
Read training status from shared file
|
|
|
|
Returns:
|
|
TrainingStatus object or None if not available
|
|
"""
|
|
try:
|
|
data = self._read_json_safe(self.training_status_file)
|
|
if not data:
|
|
return None
|
|
return TrainingStatus.from_dict(data)
|
|
except Exception as e:
|
|
logger.error(f"Failed to read training status: {e}")
|
|
return None
|
|
|
|
def write_dashboard_state(self, state: DashboardState) -> None:
|
|
"""
|
|
Write dashboard state to shared file
|
|
|
|
Args:
|
|
state: DashboardState object to write
|
|
"""
|
|
try:
|
|
data = state.to_dict()
|
|
self._write_json_atomic(self.dashboard_state_file, data)
|
|
logger.debug("Dashboard state written successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to write dashboard state: {e}")
|
|
raise
|
|
|
|
def read_dashboard_state(self) -> Optional[DashboardState]:
|
|
"""
|
|
Read dashboard state from shared file
|
|
|
|
Returns:
|
|
DashboardState object or None if not available
|
|
"""
|
|
try:
|
|
data = self._read_json_safe(self.dashboard_state_file)
|
|
if not data:
|
|
return None
|
|
return DashboardState.from_dict(data)
|
|
except Exception as e:
|
|
logger.error(f"Failed to read dashboard state: {e}")
|
|
return None
|
|
|
|
def write_process_status(self, status: ProcessStatus) -> None:
|
|
"""
|
|
Write process status to shared file
|
|
|
|
Args:
|
|
status: ProcessStatus object to write
|
|
"""
|
|
try:
|
|
data = status.to_dict()
|
|
self._write_json_atomic(self.process_status_file, data)
|
|
logger.debug("Process status written successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to write process status: {e}")
|
|
raise
|
|
|
|
def read_process_status(self) -> Optional[ProcessStatus]:
|
|
"""
|
|
Read process status from shared file
|
|
|
|
Returns:
|
|
ProcessStatus object or None if not available
|
|
"""
|
|
try:
|
|
data = self._read_json_safe(self.process_status_file)
|
|
if not data:
|
|
return None
|
|
return ProcessStatus.from_dict(data)
|
|
except Exception as e:
|
|
logger.error(f"Failed to read process status: {e}")
|
|
return None
|
|
|
|
def write_market_data(self, data: Dict[str, Any]) -> None:
|
|
"""
|
|
Write market data to shared file
|
|
|
|
Args:
|
|
data: Market data dictionary to write
|
|
"""
|
|
try:
|
|
# Add timestamp to market data
|
|
data['timestamp'] = datetime.now().isoformat()
|
|
self._write_json_atomic(self.market_data_file, data)
|
|
logger.debug("Market data written successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to write market data: {e}")
|
|
raise
|
|
|
|
def read_market_data(self) -> Dict[str, Any]:
|
|
"""
|
|
Read market data from shared file
|
|
|
|
Returns:
|
|
Dictionary containing market data
|
|
"""
|
|
try:
|
|
return self._read_json_safe(self.market_data_file)
|
|
except Exception as e:
|
|
logger.error(f"Failed to read market data: {e}")
|
|
return {}
|
|
|
|
def write_model_metrics(self, metrics: Dict[str, Any]) -> None:
|
|
"""
|
|
Write model metrics to shared file
|
|
|
|
Args:
|
|
metrics: Model metrics dictionary to write
|
|
"""
|
|
try:
|
|
# Add timestamp to metrics
|
|
metrics['timestamp'] = datetime.now().isoformat()
|
|
self._write_json_atomic(self.model_metrics_file, metrics)
|
|
logger.debug("Model metrics written successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to write model metrics: {e}")
|
|
raise
|
|
|
|
def read_model_metrics(self) -> Dict[str, Any]:
|
|
"""
|
|
Read model metrics from shared file
|
|
|
|
Returns:
|
|
Dictionary containing model metrics
|
|
"""
|
|
try:
|
|
return self._read_json_safe(self.model_metrics_file)
|
|
except Exception as e:
|
|
logger.error(f"Failed to read model metrics: {e}")
|
|
return {}
|
|
|
|
def cleanup(self) -> None:
|
|
"""
|
|
Clean up shared data files
|
|
"""
|
|
try:
|
|
for file_path in [
|
|
self.training_status_file,
|
|
self.dashboard_state_file,
|
|
self.process_status_file,
|
|
self.market_data_file,
|
|
self.model_metrics_file
|
|
]:
|
|
if file_path.exists():
|
|
file_path.unlink()
|
|
logger.debug(f"Removed {file_path}")
|
|
|
|
# Remove directory if empty
|
|
if self.data_dir.exists() and not any(self.data_dir.iterdir()):
|
|
self.data_dir.rmdir()
|
|
logger.debug(f"Removed empty directory {self.data_dir}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to cleanup shared data: {e}")
|
|
|
|
def get_data_age(self, data_type: str) -> Optional[float]:
|
|
"""
|
|
Get the age of data in seconds
|
|
|
|
Args:
|
|
data_type: Type of data ('training', 'dashboard', 'process', 'market', 'metrics')
|
|
|
|
Returns:
|
|
Age in seconds or None if file doesn't exist
|
|
"""
|
|
file_map = {
|
|
'training': self.training_status_file,
|
|
'dashboard': self.dashboard_state_file,
|
|
'process': self.process_status_file,
|
|
'market': self.market_data_file,
|
|
'metrics': self.model_metrics_file
|
|
}
|
|
|
|
file_path = file_map.get(data_type)
|
|
if not file_path or not file_path.exists():
|
|
return None
|
|
|
|
try:
|
|
mtime = file_path.stat().st_mtime
|
|
return time.time() - mtime
|
|
except Exception as e:
|
|
logger.error(f"Failed to get data age for {data_type}: {e}")
|
|
return None |