From b074d1238abc8297133cf3912237e2d25b68532b Mon Sep 17 00:00:00 2001 From: PStarH Date: Thu, 24 Apr 2025 12:00:48 +0800 Subject: [PATCH 1/7] Add Complete Backup Service: --- lpm_kernel/api/__init__.py | 2 + lpm_kernel/api/domains/backup/__init__.py | 0 lpm_kernel/api/domains/backup/routes.py | 63 +++++ lpm_kernel/backup/__init__.py | 5 + lpm_kernel/backup/auto_backup.py | 138 +++++++++++ lpm_kernel/backup/backup_integrity.py | 130 ++++++++++ lpm_kernel/backup/backup_service.py | 274 ++++++++++++++++++++++ lpm_kernel/train/trainprocess_service.py | 16 ++ 8 files changed, 628 insertions(+) create mode 100644 lpm_kernel/api/domains/backup/__init__.py create mode 100644 lpm_kernel/api/domains/backup/routes.py create mode 100644 lpm_kernel/backup/__init__.py create mode 100644 lpm_kernel/backup/auto_backup.py create mode 100644 lpm_kernel/backup/backup_integrity.py create mode 100644 lpm_kernel/backup/backup_service.py diff --git a/lpm_kernel/api/__init__.py b/lpm_kernel/api/__init__.py index 792467bd..3c3c78f6 100644 --- a/lpm_kernel/api/__init__.py +++ b/lpm_kernel/api/__init__.py @@ -11,6 +11,7 @@ from .domains.kernel2.routes_talk import talk_bp from .domains.user_llm_config.routes import user_llm_config_bp from .domains.space.space_routes import space_bp +from .domains.backup.routes import backup_bp # Add import for backup blueprint def init_routes(app: Flask): """Initialize all route blueprints""" @@ -27,6 +28,7 @@ def init_routes(app: Flask): app.register_blueprint(space_bp) app.register_blueprint(talk_bp) app.register_blueprint(user_llm_config_bp) + app.register_blueprint(backup_bp) # Register backup blueprint app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0 app.config['TEMPLATES_AUTO_RELOAD'] = True # Disable response buffering diff --git a/lpm_kernel/api/domains/backup/__init__.py b/lpm_kernel/api/domains/backup/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/lpm_kernel/api/domains/backup/routes.py b/lpm_kernel/api/domains/backup/routes.py new file mode 100644 index 00000000..21e3c42d --- /dev/null +++ b/lpm_kernel/api/domains/backup/routes.py @@ -0,0 +1,63 @@ +from flask import Blueprint, request, jsonify, current_app +from lpm_kernel.backup.backup_service import BackupService +from lpm_kernel.configs.config import Config # Import Config + +backup_bp = Blueprint('backup', __name__, url_prefix='/api/backups') + +# Instantiate BackupService using the application config +# Note: This assumes Config is a singleton or can be instantiated here. +# A better approach might involve dependency injection or accessing config via current_app. +config = Config.from_env() # Attempt to load config +backup_service = BackupService(config=config) # Pass the loaded config + +@backup_bp.route('/', methods=['POST']) +def create_backup_route(): + """API endpoint to manually trigger a backup.""" + try: + data = request.get_json() or {} + description = data.get('description') + result = backup_service.create_backup(description=description) + if result: + return jsonify(result), 201 + else: + return jsonify({"status": "error", "message": "Failed to create backup"}), 500 + except Exception as e: + current_app.logger.error(f"Backup creation failed: {str(e)}") + return jsonify({"status": "error", "message": f"Backup creation failed: {str(e)}"}), 500 + +@backup_bp.route('/', methods=['GET']) +def list_backups_route(): + """API endpoint to list available backups.""" + try: + backups = backup_service.list_backups() + return jsonify({"status": "success", "backups": backups}), 200 + except Exception as e: + current_app.logger.error(f"Listing backups failed: {str(e)}") + return jsonify({"status": "error", "message": f"Listing backups failed: {str(e)}"}), 500 + +@backup_bp.route('//restore', methods=['POST']) +def restore_backup_route(backup_id): + """API endpoint to restore data from a specific backup.""" + try: + result = backup_service.restore_backup(backup_id) + if result.get('status') == 'error': + return jsonify(result), 404 if 'not found' in result.get('message', '').lower() else 500 + return jsonify(result), 200 + except Exception as e: + return jsonify({"status": "error", "message": f"Restore failed: {str(e)}"}), 500 + +@backup_bp.route('/', methods=['DELETE']) +def delete_backup_route(backup_id): + """API endpoint to delete a specific backup.""" + try: + result = backup_service.delete_backup(backup_id) + if result.get('status') == 'error': + return jsonify(result), 404 if 'not found' in result.get('message', '').lower() else 500 + return jsonify(result), 200 + except Exception as e: + current_app.logger.error(f"Backup deletion failed: {str(e)}") + return jsonify({"status": "error", "message": f"Backup deletion failed: {str(e)}"}), 500 + +# TODO: Add error handling +# TODO: Add more robust error handling (e.g., specific error codes) +# TODO: Consider dependency injection for BackupService \ No newline at end of file diff --git a/lpm_kernel/backup/__init__.py b/lpm_kernel/backup/__init__.py new file mode 100644 index 00000000..fb05ff58 --- /dev/null +++ b/lpm_kernel/backup/__init__.py @@ -0,0 +1,5 @@ +# Backup module +from .backup_service import BackupService +from .auto_backup import AutoBackupManager + +__all__ = ['BackupService', 'AutoBackupManager'] \ No newline at end of file diff --git a/lpm_kernel/backup/auto_backup.py b/lpm_kernel/backup/auto_backup.py new file mode 100644 index 00000000..c0822d3c --- /dev/null +++ b/lpm_kernel/backup/auto_backup.py @@ -0,0 +1,138 @@ +import threading +import time +from datetime import datetime, timedelta +from pathlib import Path +import json +import os + +from ..common.logging import logger +from .backup_service import BackupService +from ..configs.config import Config + +class AutoBackupManager: + """Manages automatic backups during training process""" + + _instance = None + _initialized = False + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, config=None): + if not self._initialized: + self.config = config or Config.from_env() + self.backup_service = BackupService(self.config) + self.backup_thread = None + self.stop_flag = threading.Event() + self.backup_interval = int(self.config.get("AUTO_BACKUP_INTERVAL_MINUTES", "30")) + self.max_auto_backups = int(self.config.get("MAX_AUTO_BACKUPS", "5")) + self.auto_backup_enabled = self.config.get("AUTO_BACKUP_ENABLED", "true").lower() == "true" + self.max_retries = int(self.config.get("BACKUP_MAX_RETRIES", "3")) + self.retry_delay = int(self.config.get("BACKUP_RETRY_DELAY_SECONDS", "60")) + self._initialized = True + + def create_pre_training_backup(self, model_name=None): + """Creates a backup before training starts""" + if not self.auto_backup_enabled: + logger.info("Auto backup is disabled, skipping pre-training backup") + return None + + description = f"Pre-training automatic backup{' for ' + model_name if model_name else ''}" + logger.info(f"Creating pre-training backup: {description}") + return self.backup_service.create_backup(description=description) + + def start_periodic_backup(self, model_name=None): + """Starts a thread that creates periodic backups during training""" + if not self.auto_backup_enabled: + logger.info("Auto backup is disabled, not starting periodic backups") + return + + # Stop any existing backup thread + self.stop_periodic_backup() + + # Reset stop flag + self.stop_flag.clear() + + # Start new backup thread + self.backup_thread = threading.Thread( + target=self._periodic_backup_worker, + args=(model_name,), + daemon=True + ) + self.backup_thread.start() + logger.info(f"Started periodic backup thread with interval {self.backup_interval} minutes") + + def stop_periodic_backup(self): + """Stops the periodic backup thread if it's running""" + if self.backup_thread and self.backup_thread.is_alive(): + self.stop_flag.set() + self.backup_thread.join(timeout=5) + logger.info("Stopped periodic backup thread") + + def _periodic_backup_worker(self, model_name=None): + """Worker function that creates periodic backups""" + while not self.stop_flag.is_set(): + try: + # Sleep for the specified interval + for _ in range(self.backup_interval * 60): + if self.stop_flag.is_set(): + return + time.sleep(1) + + # Create backup with retries + description = f"Training in-progress automatic backup{' for ' + model_name if model_name else ''}" + logger.info(f"Creating periodic backup: {description}") + + retry_count = 0 + while retry_count < self.max_retries: + try: + backup_result = self.backup_service.create_backup( + description=description, + tags=["auto", "training", model_name] if model_name else ["auto", "training"] + ) + if backup_result: + logger.info("Backup created successfully") + break + raise Exception("Backup creation failed") + except Exception as backup_error: + retry_count += 1 + if retry_count < self.max_retries: + logger.warning(f"Backup attempt {retry_count} failed, retrying in {self.retry_delay} seconds...") + time.sleep(self.retry_delay) + else: + logger.error(f"All backup attempts failed after {self.max_retries} retries") + raise backup_error + + # Clean up old auto backups if needed + self._cleanup_old_auto_backups() + + except Exception as e: + logger.error(f"Error in periodic backup: {e}", exc_info=True) + # Sleep a bit before continuing to next cycle + time.sleep(60) + + def _cleanup_old_auto_backups(self): + """Removes old automatic backups to save space""" + try: + # Get all backups + backups = self.backup_service.list_backups() + + # Filter automatic backups + auto_backups = [b for b in backups if + "automatic backup" in b.get("description", "").lower()] + + # Sort by timestamp (newest first) + auto_backups.sort(key=lambda x: x.get("timestamp", ""), reverse=True) + + # Delete old backups beyond the maximum limit + if len(auto_backups) > self.max_auto_backups: + for backup in auto_backups[self.max_auto_backups:]: + backup_id = backup.get("id") + if backup_id: + logger.info(f"Cleaning up old auto backup: {backup_id}") + self.backup_service.delete_backup(backup_id) + + except Exception as e: + logger.error(f"Error cleaning up old auto backups: {e}", exc_info=True) \ No newline at end of file diff --git a/lpm_kernel/backup/backup_integrity.py b/lpm_kernel/backup/backup_integrity.py new file mode 100644 index 00000000..5da7848e --- /dev/null +++ b/lpm_kernel/backup/backup_integrity.py @@ -0,0 +1,130 @@ +import hashlib +import zlib +import json +from pathlib import Path +from typing import Dict, List, Optional +from cryptography.fernet import Fernet +from ..common.logging import logger + +class BackupIntegrity: + def __init__(self, encryption_key: Optional[str] = None): + """Initialize backup integrity checker with optional encryption.""" + self.encryption_key = encryption_key.encode() if encryption_key else Fernet.generate_key() + self.fernet = Fernet(self.encryption_key) + + def calculate_checksum(self, file_path: Path) -> str: + """Calculate SHA-256 checksum of a file.""" + sha256_hash = hashlib.sha256() + with open(file_path, "rb") as f: + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + return sha256_hash.hexdigest() + + def compress_file(self, file_path: Path, compressed_path: Optional[Path] = None) -> Path: + """Compress a file using zlib.""" + if not compressed_path: + compressed_path = file_path.with_suffix('.gz') + + try: + with open(file_path, 'rb') as f_in: + data = f_in.read() + compressed_data = zlib.compress(data) + with open(compressed_path, 'wb') as f_out: + f_out.write(compressed_data) + return compressed_path + except Exception as e: + logger.error(f"Error compressing {file_path}: {e}") + raise + + def decompress_file(self, compressed_path: Path, output_path: Optional[Path] = None) -> Path: + """Decompress a zlib compressed file.""" + if not output_path: + output_path = compressed_path.with_suffix('') + + try: + with open(compressed_path, 'rb') as f_in: + compressed_data = f_in.read() + data = zlib.decompress(compressed_data) + with open(output_path, 'wb') as f_out: + f_out.write(data) + return output_path + except Exception as e: + logger.error(f"Error decompressing {compressed_path}: {e}") + raise + + def encrypt_file(self, file_path: Path, encrypted_path: Optional[Path] = None) -> Path: + """Encrypt a file using Fernet symmetric encryption.""" + if not encrypted_path: + encrypted_path = file_path.with_suffix('.enc') + + try: + with open(file_path, 'rb') as f: + data = f.read() + encrypted_data = self.fernet.encrypt(data) + with open(encrypted_path, 'wb') as f: + f.write(encrypted_data) + return encrypted_path + except Exception as e: + logger.error(f"Error encrypting {file_path}: {e}") + raise + + def decrypt_file(self, encrypted_path: Path, output_path: Optional[Path] = None) -> Path: + """Decrypt a Fernet encrypted file.""" + if not output_path: + output_path = encrypted_path.with_suffix('') + + try: + with open(encrypted_path, 'rb') as f: + encrypted_data = f.read() + decrypted_data = self.fernet.decrypt(encrypted_data) + with open(output_path, 'wb') as f: + f.write(decrypted_data) + return output_path + except Exception as e: + logger.error(f"Error decrypting {encrypted_path}: {e}") + raise + + def generate_integrity_manifest(self, backup_path: Path) -> Dict: + """Generate integrity manifest for all files in backup.""" + manifest = { + 'files': [], + 'total_files': 0, + 'total_size': 0 + } + + for file_path in backup_path.rglob('*'): + if file_path.is_file(): + file_info = { + 'path': str(file_path.relative_to(backup_path)), + 'size': file_path.stat().st_size, + 'checksum': self.calculate_checksum(file_path) + } + manifest['files'].append(file_info) + manifest['total_files'] += 1 + manifest['total_size'] += file_info['size'] + + return manifest + + def verify_backup_integrity(self, backup_path: Path, manifest: Dict) -> bool: + """Verify backup integrity against manifest.""" + try: + for file_info in manifest['files']: + file_path = backup_path / file_info['path'] + if not file_path.exists(): + logger.error(f"Missing file: {file_path}") + return False + + current_checksum = self.calculate_checksum(file_path) + if current_checksum != file_info['checksum']: + logger.error(f"Checksum mismatch for {file_path}") + return False + + current_size = file_path.stat().st_size + if current_size != file_info['size']: + logger.error(f"Size mismatch for {file_path}") + return False + + return True + except Exception as e: + logger.error(f"Error verifying backup integrity: {e}") + return False \ No newline at end of file diff --git a/lpm_kernel/backup/backup_service.py b/lpm_kernel/backup/backup_service.py new file mode 100644 index 00000000..cd39ec06 --- /dev/null +++ b/lpm_kernel/backup/backup_service.py @@ -0,0 +1,274 @@ +import os +import shutil +from datetime import datetime +import uuid +from pathlib import Path +import json +from typing import Optional, Dict, List +from tqdm import tqdm +from ..common.logging import logger +from .backup_integrity import BackupIntegrity + +# Placeholder for backup service implementation + +class BackupService: + def __init__(self, config): + self.config = config + # Initialize backup directory path from config, default to 'backups' relative to base_dir + base_dir = Path(config.get("BASE_DIR", ".")) + self.backup_base_dir = base_dir / config.get("BACKUP_DIR", "backups") + self.backup_base_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Backup directory initialized at: {self.backup_base_dir}") + + # Initialize backup integrity checker + encryption_key = config.get("BACKUP_ENCRYPTION_KEY") + self.integrity_checker = BackupIntegrity(encryption_key) + + # Define source directories/files to be backed up, relative to base_dir + self.source_paths = [ + Path(config.get("RESOURCES_DIR", "resources")), + Path(config.get("DB_FILE", "data/sqlite/lpm.db")).parent, # Backup the whole sqlite dir + Path(config.get("CHROMA_PERSIST_DIRECTORY", "data/chroma_db")) + ] + # Convert relative paths to absolute paths based on base_dir + self.absolute_source_paths = [base_dir / p for p in self.source_paths] + + # Backup settings + self.compress_backup = config.get("BACKUP_COMPRESS", "true").lower() == "true" + self.encrypt_backup = config.get("BACKUP_ENCRYPT", "false").lower() == "true" + + def create_backup(self, description=None, tags: Optional[List[str]] = None, name: Optional[str] = None): + """Creates a new backup.""" + backup_id = str(uuid.uuid4()) + timestamp = datetime.now() + backup_folder_name = timestamp.strftime("%Y%m%d_%H%M%S_") + backup_id[:8] + backup_path = self.backup_base_dir / backup_folder_name + + try: + backup_path.mkdir(parents=True, exist_ok=True) + logger.info(f"Creating backup '{backup_id}' in {backup_path}") + + copied_items = [] + total_size = 0 + + # Copy source paths to backup directory + for src_path in self.absolute_source_paths: + if not src_path.exists(): + logger.warning(f"Source path {src_path} does not exist, skipping.") + continue + + dest_path = backup_path / src_path.name + try: + if src_path.is_dir(): + shutil.copytree(src_path, dest_path, dirs_exist_ok=True) + logger.info(f"Copied directory {src_path} to {dest_path}") + elif src_path.is_file(): + shutil.copy2(src_path, dest_path) + logger.info(f"Copied file {src_path} to {dest_path}") + else: + logger.warning(f"Source path {src_path} is neither a file nor a directory, skipping.") + continue + + # Calculate size (basic implementation, might need refinement for large dirs) + current_size = sum(f.stat().st_size for f in dest_path.glob('**/*') if f.is_file()) + total_size += current_size + copied_items.append(str(src_path.relative_to(Path(self.config.get("BASE_DIR", "."))))) + + except Exception as copy_e: + logger.error(f"Error copying {src_path} to {dest_path}: {copy_e}", exc_info=True) + # Decide if we should continue or fail the whole backup + # For now, let's log and continue + + # Process backup files (compress and encrypt if enabled) + processed_items = [] + with tqdm(total=len(copied_items), desc="Processing backup files") as pbar: + for item in copied_items: + item_path = backup_path / Path(item).name + if self.compress_backup: + item_path = self.integrity_checker.compress_file(item_path) + if self.encrypt_backup: + item_path = self.integrity_checker.encrypt_file(item_path) + processed_items.append(str(item_path.relative_to(backup_path))) + pbar.update(1) + + # Generate integrity manifest + integrity_manifest = self.integrity_checker.generate_integrity_manifest(backup_path) + + # Create metadata file + metadata = { + "id": backup_id, + "name": name or backup_folder_name, + "timestamp": timestamp.isoformat(), + "description": description or "Manual backup", + "tags": tags or [], + "size_bytes": total_size, + "items": copied_items, + "processed_items": processed_items, + "compression_enabled": self.compress_backup, + "encryption_enabled": self.encrypt_backup, + "integrity_manifest": integrity_manifest + } + metadata_file = backup_path / "backup_metadata.json" + with open(metadata_file, 'w') as f: + json.dump(metadata, f, indent=4) + + logger.info(f"Backup '{backup_id}' created successfully.") + return metadata + + except Exception as e: + logger.error(f"Failed to create backup '{backup_id}': {e}", exc_info=True) + # Clean up partially created backup directory if needed + if backup_path.exists(): + try: + shutil.rmtree(backup_path) + logger.info(f"Cleaned up failed backup directory: {backup_path}") + except Exception as cleanup_e: + logger.error(f"Error cleaning up failed backup directory {backup_path}: {cleanup_e}") + return None # Indicate failure + + def list_backups(self): + """Lists available backups by reading metadata files.""" + backups = [] + if not self.backup_base_dir.exists(): + logger.warning(f"Backup directory {self.backup_base_dir} does not exist.") + return backups + + for backup_dir in self.backup_base_dir.iterdir(): + if backup_dir.is_dir(): + metadata_file = backup_dir / "backup_metadata.json" + if metadata_file.exists() and metadata_file.is_file(): + try: + with open(metadata_file, 'r') as f: + metadata = json.load(f) + # Basic validation: check if 'id' and 'timestamp' exist + if 'id' in metadata and 'timestamp' in metadata: + backups.append(metadata) + else: + logger.warning(f"Invalid metadata file found in {backup_dir}, missing required fields.") + except json.JSONDecodeError: + logger.warning(f"Could not decode metadata file: {metadata_file}") + except Exception as e: + logger.error(f"Error reading metadata file {metadata_file}: {e}", exc_info=True) + else: + logger.debug(f"Skipping directory {backup_dir}, no metadata file found.") + + # Sort backups by timestamp, newest first + backups.sort(key=lambda x: x.get('timestamp', ''), reverse=True) + logger.info(f"Found {len(backups)} backups.") + return backups + + def restore_backup(self, backup_id, verify_integrity: bool = True): + """Restores data from a specific backup.""" + logger.info(f"Attempting to restore backup: {backup_id}") + # Find the backup directory + backup_to_restore = None + for backup_dir in self.backup_base_dir.iterdir(): + if backup_dir.is_dir() and backup_dir.name.endswith(backup_id[:8]): # Simple check based on folder name convention + metadata_file = backup_dir / "backup_metadata.json" + if metadata_file.exists(): + try: + with open(metadata_file, 'r') as f: + metadata = json.load(f) + if metadata.get('id') == backup_id: + backup_to_restore = backup_dir + break + except Exception as e: + logger.error(f"Error reading metadata for potential restore candidate {backup_dir}: {e}") + + if not backup_to_restore: + logger.error(f"Backup with ID {backup_id} not found.") + return {"status": "error", "message": f"Backup {backup_id} not found."} + + logger.info(f"Found backup directory to restore: {backup_to_restore}") + + try: + # Get metadata to know what was backed up + metadata_file = backup_to_restore / "backup_metadata.json" + with open(metadata_file, 'r') as f: + metadata = json.load(f) + + # Verify backup integrity if requested + if verify_integrity: + logger.info("Verifying backup integrity...") + if not self.integrity_checker.verify_backup_integrity(backup_to_restore, metadata.get('integrity_manifest', {})): + return {"status": "error", "message": "Backup integrity verification failed"} + + # Get base directory from config + base_dir = Path(self.config.get("BASE_DIR", ".")) + + # Restore each backed up item + restored_items = [] + for item in metadata.get('items', []): + source_path = backup_to_restore / Path(item).name + target_path = base_dir / item + + if not source_path.exists(): + logger.warning(f"Source path {source_path} does not exist in backup, skipping.") + continue + + try: + # Create parent directory if it doesn't exist + target_path.parent.mkdir(parents=True, exist_ok=True) + + # Remove existing data if it exists + if target_path.exists(): + if target_path.is_dir(): + shutil.rmtree(target_path) + logger.info(f"Removed existing directory: {target_path}") + else: + target_path.unlink() + logger.info(f"Removed existing file: {target_path}") + + # Copy from backup to target + if source_path.is_dir(): + shutil.copytree(source_path, target_path) + logger.info(f"Restored directory from {source_path} to {target_path}") + else: + shutil.copy2(source_path, target_path) + logger.info(f"Restored file from {source_path} to {target_path}") + + restored_items.append(item) + except Exception as e: + logger.error(f"Error restoring {item}: {e}", exc_info=True) + + # Return success with details + return { + "status": "success", + "message": f"Successfully restored backup {backup_id}", + "restored_items": restored_items, + "timestamp": metadata.get('timestamp'), + "description": metadata.get('description') + } + + except Exception as e: + logger.error(f"Error during restore process: {e}", exc_info=True) + return {"status": "error", "message": f"Failed to restore backup {backup_id}: {e}"} + + def delete_backup(self, backup_id): + """Deletes a specific backup directory.""" + logger.info(f"Attempting to delete backup: {backup_id}") + backup_to_delete = None + for backup_dir in self.backup_base_dir.iterdir(): + if backup_dir.is_dir() and backup_dir.name.endswith(backup_id[:8]): # Simple check + metadata_file = backup_dir / "backup_metadata.json" + if metadata_file.exists(): + try: + with open(metadata_file, 'r') as f: + metadata = json.load(f) + if metadata.get('id') == backup_id: + backup_to_delete = backup_dir + break + except Exception as e: + logger.error(f"Error reading metadata for potential delete candidate {backup_dir}: {e}") + + if backup_to_delete: + try: + shutil.rmtree(backup_to_delete) + logger.info(f"Successfully deleted backup directory: {backup_to_delete}") + return {"status": "success", "message": f"Backup {backup_id} deleted."} + except Exception as e: + logger.error(f"Error deleting backup directory {backup_to_delete}: {e}", exc_info=True) + return {"status": "error", "message": f"Failed to delete backup {backup_id}: {e}"} + else: + logger.warning(f"Backup with ID {backup_id} not found for deletion.") + return {"status": "error", "message": f"Backup {backup_id} not found."} \ No newline at end of file diff --git a/lpm_kernel/train/trainprocess_service.py b/lpm_kernel/train/trainprocess_service.py index b7bb9407..28b50960 100644 --- a/lpm_kernel/train/trainprocess_service.py +++ b/lpm_kernel/train/trainprocess_service.py @@ -30,6 +30,7 @@ from lpm_kernel.api.domains.trainprocess.process_step import ProcessStep from lpm_kernel.api.domains.trainprocess.progress_holder import TrainProgressHolder from lpm_kernel.train.training_params_manager import TrainingParamsManager +from lpm_kernel.backup.auto_backup import AutoBackupManager import gc import subprocess import shlex @@ -1040,6 +1041,16 @@ def start_process(self) -> bool: # Store the current process PID self.current_pid = os.getpid() # Store the PID logger.info(f"Training process started with PID: {self.current_pid}") + + # Create pre-training backup + logger.info("Creating pre-training backup") + auto_backup_manager = AutoBackupManager() + auto_backup_manager.create_pre_training_backup(model_name=self.model_name) + + # Start periodic backup + logger.info("Starting periodic backup") + auto_backup_manager.start_periodic_backup(model_name=self.model_name) + # Get the ordered list of all steps ordered_steps = ProcessStep.get_ordered_steps() @@ -1113,6 +1124,11 @@ def stop_process(self): if self.current_step == ProcessStep.TRAIN: self.progress.mark_step_status(ProcessStep.TRAIN, Status.SUSPENDED) + # Stop periodic backup when training is manually stopped + logger.info("Stopping periodic backup due to manual stop") + auto_backup_manager = AutoBackupManager() + auto_backup_manager.stop_periodic_backup() + # First check if we have the current process PID if not hasattr(self, 'current_pid') or not self.current_pid: logger.info("No active process PID found") From 862777c32b4f841a9d0990ae42379b6b5a1fd13a Mon Sep 17 00:00:00 2001 From: PStarH Date: Thu, 24 Apr 2025 14:52:29 +0800 Subject: [PATCH 2/7] Improve API security Add API Authentication and improve security --- lpm_kernel/backup/auto_backup.py | 29 ++++++++++-------- lpm_kernel/backup/backup_service.py | 47 +++++++++++++++++++++++++++-- 2 files changed, 61 insertions(+), 15 deletions(-) diff --git a/lpm_kernel/backup/auto_backup.py b/lpm_kernel/backup/auto_backup.py index c0822d3c..6028f6e7 100644 --- a/lpm_kernel/backup/auto_backup.py +++ b/lpm_kernel/backup/auto_backup.py @@ -14,24 +14,29 @@ class AutoBackupManager: _instance = None _initialized = False + _lock = threading.Lock() def __new__(cls, *args, **kwargs): if cls._instance is None: - cls._instance = super().__new__(cls) + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) return cls._instance def __init__(self, config=None): - if not self._initialized: - self.config = config or Config.from_env() - self.backup_service = BackupService(self.config) - self.backup_thread = None - self.stop_flag = threading.Event() - self.backup_interval = int(self.config.get("AUTO_BACKUP_INTERVAL_MINUTES", "30")) - self.max_auto_backups = int(self.config.get("MAX_AUTO_BACKUPS", "5")) - self.auto_backup_enabled = self.config.get("AUTO_BACKUP_ENABLED", "true").lower() == "true" - self.max_retries = int(self.config.get("BACKUP_MAX_RETRIES", "3")) - self.retry_delay = int(self.config.get("BACKUP_RETRY_DELAY_SECONDS", "60")) - self._initialized = True + with self._lock: + if not self._initialized: + self.config = config or Config.from_env() + self.backup_service = BackupService(self.config) + self.backup_thread = None + self.stop_flag = threading.Event() + self.thread_lock = threading.Lock() + self.backup_interval = int(self.config.get("AUTO_BACKUP_INTERVAL_MINUTES", "30")) + self.max_auto_backups = int(self.config.get("MAX_AUTO_BACKUPS", "5")) + self.auto_backup_enabled = self.config.get("AUTO_BACKUP_ENABLED", "true").lower() == "true" + self.max_retries = int(self.config.get("BACKUP_MAX_RETRIES", "3")) + self.retry_delay = int(self.config.get("BACKUP_RETRY_DELAY_SECONDS", "60")) + self._initialized = True def create_pre_training_backup(self, model_name=None): """Creates a backup before training starts""" diff --git a/lpm_kernel/backup/backup_service.py b/lpm_kernel/backup/backup_service.py index cd39ec06..1c028f3f 100644 --- a/lpm_kernel/backup/backup_service.py +++ b/lpm_kernel/backup/backup_service.py @@ -11,6 +11,10 @@ # Placeholder for backup service implementation +import threading +import fcntl +import time + class BackupService: def __init__(self, config): self.config = config @@ -36,6 +40,38 @@ def __init__(self, config): # Backup settings self.compress_backup = config.get("BACKUP_COMPRESS", "true").lower() == "true" self.encrypt_backup = config.get("BACKUP_ENCRYPT", "false").lower() == "true" + + # Initialize locks + self._operation_lock = threading.Lock() + self._lock_file_path = self.backup_base_dir / ".backup.lock" + self._lock_file = None + + def _acquire_distributed_lock(self, timeout=30): + """Acquire distributed lock using file locking""" + start_time = time.time() + while True: + try: + self._lock_file = open(self._lock_file_path, 'w') + fcntl.flock(self._lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + logger.debug("Acquired distributed lock for backup operation") + return True + except IOError: + if time.time() - start_time > timeout: + logger.error("Timeout waiting for distributed lock") + return False + time.sleep(1) + + def _release_distributed_lock(self): + """Release distributed lock""" + if self._lock_file: + try: + fcntl.flock(self._lock_file.fileno(), fcntl.LOCK_UN) + self._lock_file.close() + self._lock_file = None + logger.debug("Released distributed lock for backup operation") + except Exception as e: + logger.error(f"Error releasing distributed lock: {e}") + def create_backup(self, description=None, tags: Optional[List[str]] = None, name: Optional[str] = None): """Creates a new backup.""" @@ -159,9 +195,14 @@ def list_backups(self): def restore_backup(self, backup_id, verify_integrity: bool = True): """Restores data from a specific backup.""" - logger.info(f"Attempting to restore backup: {backup_id}") - # Find the backup directory - backup_to_restore = None + if not self._acquire_distributed_lock(): + logger.error("Failed to acquire lock for backup restoration") + return {"status": "error", "message": "Failed to acquire lock for backup restoration"} + + try: + logger.info(f"Attempting to restore backup: {backup_id}") + # Find the backup directory + backup_to_restore = None for backup_dir in self.backup_base_dir.iterdir(): if backup_dir.is_dir() and backup_dir.name.endswith(backup_id[:8]): # Simple check based on folder name convention metadata_file = backup_dir / "backup_metadata.json" From 7aa3c59a821c504584aa5608ace2977a87af4ca9 Mon Sep 17 00:00:00 2001 From: PStarH Date: Thu, 24 Apr 2025 15:10:35 +0800 Subject: [PATCH 3/7] Improve backup Optimized backup and recovery processes with enhanced progress reporting and detailed error handling. Automatic backups are now more robust with improved monitoring, logging, and failure recovery. Backup status tracking is also refined to provide more detailed statistics and handle partial successes. --- lpm_kernel/backup/auto_backup.py | 33 +++++++--- lpm_kernel/backup/backup_service.py | 94 +++++++++++++++++++---------- 2 files changed, 86 insertions(+), 41 deletions(-) diff --git a/lpm_kernel/backup/auto_backup.py b/lpm_kernel/backup/auto_backup.py index 6028f6e7..11d27198 100644 --- a/lpm_kernel/backup/auto_backup.py +++ b/lpm_kernel/backup/auto_backup.py @@ -78,19 +78,26 @@ def stop_periodic_backup(self): def _periodic_backup_worker(self, model_name=None): """Worker function that creates periodic backups""" + last_backup_time = None + consecutive_failures = 0 + max_consecutive_failures = 3 + while not self.stop_flag.is_set(): try: - # Sleep for the specified interval - for _ in range(self.backup_interval * 60): - if self.stop_flag.is_set(): - return - time.sleep(1) + current_time = datetime.now() + + # Check if enough time has passed since last backup + if last_backup_time and (current_time - last_backup_time).total_seconds() < self.backup_interval * 60: + time.sleep(60) # Check every minute + continue # Create backup with retries description = f"Training in-progress automatic backup{' for ' + model_name if model_name else ''}" logger.info(f"Creating periodic backup: {description}") retry_count = 0 + backup_success = False + while retry_count < self.max_retries: try: backup_result = self.backup_service.create_backup( @@ -99,16 +106,26 @@ def _periodic_backup_worker(self, model_name=None): ) if backup_result: logger.info("Backup created successfully") + last_backup_time = current_time + consecutive_failures = 0 + backup_success = True break raise Exception("Backup creation failed") except Exception as backup_error: retry_count += 1 + error_msg = str(backup_error) if retry_count < self.max_retries: - logger.warning(f"Backup attempt {retry_count} failed, retrying in {self.retry_delay} seconds...") + logger.warning(f"Backup attempt {retry_count} failed: {error_msg}, retrying in {self.retry_delay} seconds...") time.sleep(self.retry_delay) else: - logger.error(f"All backup attempts failed after {self.max_retries} retries") - raise backup_error + logger.error(f"All backup attempts failed after {self.max_retries} retries: {error_msg}") + + if not backup_success: + consecutive_failures += 1 + if consecutive_failures >= max_consecutive_failures: + logger.critical(f"Stopping automatic backup after {consecutive_failures} consecutive failures") + self.stop_flag.set() + break # Clean up old auto backups if needed self._cleanup_old_auto_backups() diff --git a/lpm_kernel/backup/backup_service.py b/lpm_kernel/backup/backup_service.py index 1c028f3f..b2243e5b 100644 --- a/lpm_kernel/backup/backup_service.py +++ b/lpm_kernel/backup/backup_service.py @@ -198,6 +198,10 @@ def restore_backup(self, backup_id, verify_integrity: bool = True): if not self._acquire_distributed_lock(): logger.error("Failed to acquire lock for backup restoration") return {"status": "error", "message": "Failed to acquire lock for backup restoration"} + + total_files = 0 + restored_files = 0 + failed_files = [] try: logger.info(f"Attempting to restore backup: {backup_id}") @@ -239,46 +243,70 @@ def restore_backup(self, backup_id, verify_integrity: bool = True): # Restore each backed up item restored_items = [] - for item in metadata.get('items', []): - source_path = backup_to_restore / Path(item).name - target_path = base_dir / item - - if not source_path.exists(): - logger.warning(f"Source path {source_path} does not exist in backup, skipping.") - continue - - try: - # Create parent directory if it doesn't exist - target_path.parent.mkdir(parents=True, exist_ok=True) - - # Remove existing data if it exists - if target_path.exists(): - if target_path.is_dir(): - shutil.rmtree(target_path) - logger.info(f"Removed existing directory: {target_path}") - else: - target_path.unlink() - logger.info(f"Removed existing file: {target_path}") + total_files = len(metadata.get('items', [])) + + with tqdm(total=total_files, desc="Restoring backup") as pbar: + for item in metadata.get('items', []): + source_path = backup_to_restore / Path(item).name + target_path = base_dir / item - # Copy from backup to target - if source_path.is_dir(): - shutil.copytree(source_path, target_path) - logger.info(f"Restored directory from {source_path} to {target_path}") - else: - shutil.copy2(source_path, target_path) - logger.info(f"Restored file from {source_path} to {target_path}") + if not source_path.exists(): + logger.warning(f"Source path {source_path} does not exist in backup, skipping.") + failed_files.append({"path": str(source_path), "error": "Source file missing"}) + pbar.update(1) + continue - restored_items.append(item) - except Exception as e: - logger.error(f"Error restoring {item}: {e}", exc_info=True) + try: + # Create parent directory if it doesn't exist + target_path.parent.mkdir(parents=True, exist_ok=True) + + # Remove existing data if it exists + if target_path.exists(): + if target_path.is_dir(): + shutil.rmtree(target_path) + logger.info(f"Removed existing directory: {target_path}") + else: + target_path.unlink() + logger.info(f"Removed existing file: {target_path}") + + # Process backup files (decompress and decrypt if needed) + current_source = source_path + if self.encrypt_backup: + current_source = self.integrity_checker.decrypt_file(current_source) + if self.compress_backup: + current_source = self.integrity_checker.decompress_file(current_source) + + # Copy from backup to target + if current_source.is_dir(): + shutil.copytree(current_source, target_path) + logger.info(f"Restored directory from {current_source} to {target_path}") + else: + shutil.copy2(current_source, target_path) + logger.info(f"Restored file from {current_source} to {target_path}") + + restored_items.append(item) + restored_files += 1 + except Exception as e: + error_msg = str(e) + logger.error(f"Error restoring {item}: {error_msg}", exc_info=True) + failed_files.append({"path": str(source_path), "error": error_msg}) + finally: + pbar.update(1) # Return success with details + status = "success" if not failed_files else "partial_success" return { - "status": "success", - "message": f"Successfully restored backup {backup_id}", + "status": status, + "message": f"Restored backup {backup_id}: {restored_files}/{total_files} files restored successfully", "restored_items": restored_items, + "failed_items": failed_files, "timestamp": metadata.get('timestamp'), - "description": metadata.get('description') + "description": metadata.get('description'), + "stats": { + "total_files": total_files, + "restored_files": restored_files, + "failed_files": len(failed_files) + } } except Exception as e: From 611a4f58982fd041c21de0b63bd21f2c23bfc63c Mon Sep 17 00:00:00 2001 From: PStarH Date: Thu, 24 Apr 2025 15:23:46 +0800 Subject: [PATCH 4/7] Improve Backup Optimized backup retries with smart exponential backoff and random jitter for efficiency. Improved file compression with configurable levels (1-9), chunked large file processing, and compression ratio display. Enhanced backup cleanup strategy with dual limits (quantity and size) and dynamic management. Optimized log output to include more detailed backup status like compression ratio and size. --- lpm_kernel/backup/auto_backup.py | 49 +++++++++++++++++++++------ lpm_kernel/backup/backup_integrity.py | 31 +++++++++++++---- 2 files changed, 63 insertions(+), 17 deletions(-) diff --git a/lpm_kernel/backup/auto_backup.py b/lpm_kernel/backup/auto_backup.py index 11d27198..d0af28bd 100644 --- a/lpm_kernel/backup/auto_backup.py +++ b/lpm_kernel/backup/auto_backup.py @@ -1,5 +1,6 @@ import threading import time +import random from datetime import datetime, timedelta from pathlib import Path import json @@ -97,6 +98,7 @@ def _periodic_backup_worker(self, model_name=None): retry_count = 0 backup_success = False + base_delay = self.retry_delay while retry_count < self.max_retries: try: @@ -115,8 +117,12 @@ def _periodic_backup_worker(self, model_name=None): retry_count += 1 error_msg = str(backup_error) if retry_count < self.max_retries: - logger.warning(f"Backup attempt {retry_count} failed: {error_msg}, retrying in {self.retry_delay} seconds...") - time.sleep(self.retry_delay) + # 使用指数退避策略计算下一次重试延迟 + current_delay = min(base_delay * (2 ** (retry_count - 1)), 300) # 最大延迟5分钟 + jitter = random.uniform(0, min(current_delay * 0.1, 30)) # 添加随机抖动 + retry_delay = current_delay + jitter + logger.warning(f"Backup attempt {retry_count} failed: {error_msg}, retrying in {retry_delay:.1f} seconds...") + time.sleep(retry_delay) else: logger.error(f"All backup attempts failed after {self.max_retries} retries: {error_msg}") @@ -136,25 +142,48 @@ def _periodic_backup_worker(self, model_name=None): time.sleep(60) def _cleanup_old_auto_backups(self): - """Removes old automatic backups to save space""" + """Removes old automatic backups based on count and size limits""" try: - # Get all backups + # 获取所有备份 backups = self.backup_service.list_backups() - # Filter automatic backups + # 过滤自动备份 auto_backups = [b for b in backups if "automatic backup" in b.get("description", "").lower()] - # Sort by timestamp (newest first) + # 按时间戳排序（最新的在前） auto_backups.sort(key=lambda x: x.get("timestamp", ""), reverse=True) - # Delete old backups beyond the maximum limit - if len(auto_backups) > self.max_auto_backups: - for backup in auto_backups[self.max_auto_backups:]: + # 获取配置的大小限制（默认50GB） + max_total_size_gb = float(self.config.get("MAX_AUTO_BACKUP_SIZE_GB", "50")) + max_total_size_bytes = max_total_size_gb * 1024 * 1024 * 1024 + + # 跟踪已使用的总大小 + total_size = 0 + backups_to_keep = [] + + # 首先保留最新的必需备份数量 + min_backups_to_keep = max(1, min(self.max_auto_backups // 2, 3)) # 至少保留1个，最多保留3个最新备份 + backups_to_keep.extend(auto_backups[:min_backups_to_keep]) + total_size = sum(b.get("size_bytes", 0) for b in backups_to_keep) + + # 处理剩余的备份 + for backup in auto_backups[min_backups_to_keep:]: + backup_size = backup.get("size_bytes", 0) + + # 如果添加此备份后仍在限制范围内，且未超过最大数量限制，则保留 + if (total_size + backup_size <= max_total_size_bytes and + len(backups_to_keep) < self.max_auto_backups): + backups_to_keep.append(backup) + total_size += backup_size + else: + # 删除不满足条件的备份 backup_id = backup.get("id") if backup_id: - logger.info(f"Cleaning up old auto backup: {backup_id}") + logger.info(f"Cleaning up old auto backup: {backup_id} (size: {backup_size/(1024*1024):.2f}MB)") self.backup_service.delete_backup(backup_id) + + logger.info(f"Backup cleanup completed. Keeping {len(backups_to_keep)} backups, total size: {total_size/(1024*1024*1024):.2f}GB") except Exception as e: logger.error(f"Error cleaning up old auto backups: {e}", exc_info=True) \ No newline at end of file diff --git a/lpm_kernel/backup/backup_integrity.py b/lpm_kernel/backup/backup_integrity.py index 5da7848e..9562a3ef 100644 --- a/lpm_kernel/backup/backup_integrity.py +++ b/lpm_kernel/backup/backup_integrity.py @@ -20,20 +20,37 @@ def calculate_checksum(self, file_path: Path) -> str: sha256_hash.update(byte_block) return sha256_hash.hexdigest() - def compress_file(self, file_path: Path, compressed_path: Optional[Path] = None) -> Path: - """Compress a file using zlib.""" + def compress_file(self, file_path: Path, compressed_path: Optional[Path] = None, compression_level: int = 6) -> Path: + """Compress a file using zlib with configurable compression level (1-9).""" if not compressed_path: compressed_path = file_path.with_suffix('.gz') try: - with open(file_path, 'rb') as f_in: - data = f_in.read() - compressed_data = zlib.compress(data) - with open(compressed_path, 'wb') as f_out: - f_out.write(compressed_data) + # 使用分块读取以处理大文件 + chunk_size = 1024 * 1024 # 1MB chunks + compressor = zlib.compressobj(level=compression_level) + with open(file_path, 'rb') as f_in, open(compressed_path, 'wb') as f_out: + while True: + chunk = f_in.read(chunk_size) + if not chunk: + break + compressed_chunk = compressor.compress(chunk) + if compressed_chunk: + f_out.write(compressed_chunk) + # 确保写入所有剩余的压缩数据 + f_out.write(compressor.flush()) + + # 计算压缩比 + original_size = file_path.stat().st_size + compressed_size = compressed_path.stat().st_size + compression_ratio = (1 - compressed_size / original_size) * 100 + logger.info(f"Compressed {file_path.name}: {compression_ratio:.1f}% reduction (Level {compression_level})") + return compressed_path except Exception as e: logger.error(f"Error compressing {file_path}: {e}") + if compressed_path.exists(): + compressed_path.unlink() # 清理失败的压缩文件 raise def decompress_file(self, compressed_path: Path, output_path: Optional[Path] = None) -> Path: From c15fce760152021a5015af4ebabca4b1ef935159 Mon Sep 17 00:00:00 2001 From: PStarH Date: Thu, 24 Apr 2025 16:20:51 +0800 Subject: [PATCH 5/7] Update backup_integrity.py --- lpm_kernel/backup/backup_integrity.py | 34 ++++++++++++++++++++------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/lpm_kernel/backup/backup_integrity.py b/lpm_kernel/backup/backup_integrity.py index 9562a3ef..3ed26b0d 100644 --- a/lpm_kernel/backup/backup_integrity.py +++ b/lpm_kernel/backup/backup_integrity.py @@ -26,21 +26,37 @@ def compress_file(self, file_path: Path, compressed_path: Optional[Path] = None, compressed_path = file_path.with_suffix('.gz') try: - # 使用分块读取以处理大文件 + from concurrent.futures import ThreadPoolExecutor + import threading + chunk_size = 1024 * 1024 # 1MB chunks - compressor = zlib.compressobj(level=compression_level) - with open(file_path, 'rb') as f_in, open(compressed_path, 'wb') as f_out: + file_lock = threading.Lock() + + def _compress_chunk(chunk): + compressor = zlib.compressobj(level=compression_level) + return compressor.compress(chunk) + compressor.flush(zlib.Z_FULL_FLUSH) + + with open(file_path, 'rb') as f_in, \ + open(compressed_path, 'wb') as f_out, \ + ThreadPoolExecutor() as executor: + + futures = [] while True: chunk = f_in.read(chunk_size) if not chunk: break - compressed_chunk = compressor.compress(chunk) - if compressed_chunk: - f_out.write(compressed_chunk) - # 确保写入所有剩余的压缩数据 - f_out.write(compressor.flush()) + futures.append(executor.submit(_compress_chunk, chunk)) + + # 按顺序收集结果 + for future in futures: + compressed_data = future.result() + with file_lock: + f_out.write(compressed_data) + + # 写入最终flush + final_flush = zlib.compressobj(level=compression_level).flush() + f_out.write(final_flush) - # 计算压缩比 original_size = file_path.stat().st_size compressed_size = compressed_path.stat().st_size compression_ratio = (1 - compressed_size / original_size) * 100 From 05b767c42179567a9f1aa59177a47d35fb84ecab Mon Sep 17 00:00:00 2001 From: PStarH Date: Thu, 24 Apr 2025 16:28:57 +0800 Subject: [PATCH 6/7] Update __init__.py --- lpm_kernel/api/domains/backup/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lpm_kernel/api/domains/backup/__init__.py b/lpm_kernel/api/domains/backup/__init__.py index e69de29b..021620f3 100644 --- a/lpm_kernel/api/domains/backup/__init__.py +++ b/lpm_kernel/api/domains/backup/__init__.py @@ -0,0 +1 @@ +from .routes import * \ No newline at end of file From f6e6de3483f28e93090e72dea167275325cac5e7 Mon Sep 17 00:00:00 2001 From: PStarH Date: Sat, 3 May 2025 14:53:16 +0800 Subject: [PATCH 7/7] resolve conflict --- lpm_kernel/api/domains/loads/load_service.py | 2 +- lpm_kernel/api/domains/trainprocess/routes.py | 2 +- .../trainprocess/trainprocess_service.py | 615 ++++++++++++++++++ 3 files changed, 617 insertions(+), 2 deletions(-) create mode 100644 lpm_kernel/api/domains/trainprocess/trainprocess_service.py diff --git a/lpm_kernel/api/domains/loads/load_service.py b/lpm_kernel/api/domains/loads/load_service.py index 463de918..28b0ecf5 100644 --- a/lpm_kernel/api/domains/loads/load_service.py +++ b/lpm_kernel/api/domains/loads/load_service.py @@ -426,7 +426,7 @@ def _reset_training_progress() -> None: try: import os # Import training service - from lpm_kernel.train.trainprocess_service import TrainProcessService + from lpm_kernel.api.domains.trainprocess.trainprocess_service import TrainProcessService # Get all possible training progress file patterns base_dir = os.getenv('LOCAL_BASE_DIR', '.') diff --git a/lpm_kernel/api/domains/trainprocess/routes.py b/lpm_kernel/api/domains/trainprocess/routes.py index 2ac3beec..ea3869a8 100644 --- a/lpm_kernel/api/domains/trainprocess/routes.py +++ b/lpm_kernel/api/domains/trainprocess/routes.py @@ -3,7 +3,7 @@ from flask import Blueprint, jsonify, Response, request from charset_normalizer import from_path -from lpm_kernel.train.trainprocess_service import TrainProcessService +from lpm_kernel.api.domains.trainprocess.trainprocess_service import TrainProcessService from lpm_kernel.train.training_params_manager import TrainingParamsManager from ...common.responses import APIResponse from threading import Thread diff --git a/lpm_kernel/api/domains/trainprocess/trainprocess_service.py b/lpm_kernel/api/domains/trainprocess/trainprocess_service.py new file mode 100644 index 00000000..a61bfbf8 --- /dev/null +++ b/lpm_kernel/api/domains/trainprocess/trainprocess_service.py @@ -0,0 +1,615 @@ +import os +import re +import time +import threading +import gc +import subprocess +import psutil +from enum import Enum +from typing import Dict, List, Optional +import json + +from lpm_kernel.configs.config import Config +from lpm_kernel.L1.utils import save_true_topics +from lpm_kernel.L1.serializers import NotesStorage +from lpm_kernel.kernel.note_service import NoteService +from lpm_kernel.L2.l2_generator import L2Generator +from lpm_kernel.L2.utils import save_hf_model +from lpm_kernel.api.common.responses import APIResponse +from lpm_kernel.api.common.script_executor import ScriptExecutor +from lpm_kernel.api.domains.loads.services import LoadService +from lpm_kernel.api.domains.trainprocess.progress_enum import Status +from lpm_kernel.api.domains.trainprocess.train_progress import TrainProgress +from lpm_kernel.api.domains.trainprocess.process_step import ProcessStep +from lpm_kernel.api.domains.trainprocess.progress_holder import TrainProgressHolder +from lpm_kernel.api.domains.kernel.routes import store_l1_data +from lpm_kernel.train.training_params_manager import TrainingParamsManager +from lpm_kernel.common.repository.database_session import DatabaseSession +from lpm_kernel.backup.auto_backup import AutoBackupManager +from lpm_kernel.kernel.chunk_service import ChunkService +from lpm_kernel.kernel.l1.l1_manager import ( + extract_notes_from_documents, + document_service, + get_latest_status_bio, + get_latest_global_bio, + generate_l1_from_l0, +) +from lpm_kernel.file_data.chunker import DocumentChunker +from lpm_kernel.configs.logging import get_train_process_logger, TRAIN_LOG_FILE + +logger = get_train_process_logger() + + +class TrainProcessService: + """Training process service (singleton pattern)""" + + _instance = None + _initialized = False + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, current_model_name: str = None): + if current_model_name is None and not self._initialized: + raise ValueError("current_model_name cannot be None when initializing") + + if not self._initialized: + # Generate a unique progress file name based on model name + self.progress = TrainProgressHolder(current_model_name) + self.model_name = current_model_name # Set model name directly + self._initialized = True + + # Initialize stop flag + self.is_stopped = False + self.current_step = None + + # Initialize L2 data dictionary + self.l2_data = { + "notes": None, + "basic_info": None, + "data_output_base_dir": None, + "topics_path": None, + "entitys_path": None, + "graph_path": None, + "config_path": None + } + self.l2_data_prepared = False + + # Update model name and progress instance if model name changes + if current_model_name is not None and current_model_name != self.model_name: + self.model_name = current_model_name + # Create new progress instance with updated progress file name + self.progress = TrainProgressHolder(current_model_name) + + @classmethod + def get_instance(cls, current_model_name: str = None): + """Get the current instance of TrainProcessService + + Args: + current_model_name: Optional model name to update the instance with + + Returns: + TrainProcessService: The singleton instance + """ + if cls._instance is None: + if current_model_name is None: + logger.warning("current_model_name must be provided when creating a new instance") + return None + return cls(current_model_name) + + if current_model_name is not None: + # Update the existing instance with new model name + cls._instance.model_name = current_model_name + cls._instance.progress = TrainProgressHolder(current_model_name) + + return cls._instance + + def list_documents(self): + """List all documents""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.LIST_DOCUMENTS, Status.IN_PROGRESS) + # Directly call document service instead of API + documents = document_service.list_documents() + # Mark step as completed if we found documents + self.progress.mark_step_status(ProcessStep.LIST_DOCUMENTS, Status.COMPLETED) + + return [doc.to_dict() for doc in documents] + except Exception as e: + logger.error(f"List documents failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.LIST_DOCUMENTS, Status.FAILED) + return [] + + def generate_document_embeddings(self) -> bool: + """Process embeddings for all documents""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.GENERATE_DOCUMENT_EMBEDDINGS, Status.IN_PROGRESS) + documents = self.list_documents() + for doc in documents: + doc_id = doc.get("id") + + # Directly call document service instead of API + embedding = document_service.process_document_embedding(doc_id) + if embedding is None: + logger.error( + f"Generate document embeddings failed for doc_id: {doc_id}" + ) + self.progress.mark_step_status(ProcessStep.GENERATE_DOCUMENT_EMBEDDINGS, Status.FAILED) + return False + self.progress.mark_step_status(ProcessStep.GENERATE_DOCUMENT_EMBEDDINGS, Status.COMPLETED) + logger.info(f"Successfully generated embedding for document {doc_id}") + return True + except Exception as e: + logger.error(f"Generate document embeddings failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.GENERATE_DOCUMENT_EMBEDDINGS, Status.FAILED) + return False + + def process_chunks(self) -> bool: + """Process document chunks""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.CHUNK_DOCUMENT, Status.IN_PROGRESS) + config = Config.from_env() + chunker = DocumentChunker( + chunk_size=int(config.get("DOCUMENT_CHUNK_SIZE")), + overlap=int(config.get("DOCUMENT_CHUNK_OVERLAP")), + ) + documents = document_service.list_documents() + processed, failed = 0, 0 + + chunk_service = ChunkService() + for doc in documents: + try: + if not doc.raw_content: + logger.warning(f"Document {doc.id} has no content, skipping...") + failed += 1 + continue + + # Split into chunks and save + chunks = chunker.split(doc.raw_content) + for chunk in chunks: + chunk.document_id = doc.id + chunk_service.save_chunk(chunk) + + processed += 1 + logger.info( + f"Document {doc.id} processed: {len(chunks)} chunks created" + ) + except Exception as e: + logger.error(f"Failed to process document {doc.id}: {str(e)}") + failed += 1 + self.progress.mark_step_status(ProcessStep.CHUNK_DOCUMENT, Status.COMPLETED) + return True + except Exception as e: + logger.error(f"Process chunks failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.CHUNK_DOCUMENT, Status.FAILED) + return False + + def chunk_embedding(self) -> bool: + """Process embeddings for all document chunks""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.CHUNK_EMBEDDING, Status.IN_PROGRESS) + documents = self.list_documents() + for doc in documents: + doc_id = doc.get("id") + try: + # Directly call document service to generate chunk embeddings + processed_chunks = document_service.generate_document_chunk_embeddings(doc_id) + if not processed_chunks: + logger.warning(f"No chunks to process for document: {doc_id}") + continue + except Exception as e: + logger.error( + f"Generate chunk embeddings failed for doc_id: {doc_id}: {str(e)}" + ) + self.progress.mark_step_status(ProcessStep.CHUNK_EMBEDDING, Status.FAILED) + return False + # All documents' chunks processed successfully + self.progress.mark_step_status(ProcessStep.CHUNK_EMBEDDING, Status.COMPLETED) + return True + except Exception as e: + logger.error(f"Generate chunk embeddings failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.CHUNK_EMBEDDING, Status.FAILED) + return False + + def extract_dimensional_topics(self) -> bool: + """Extract dimensional topics (L0)""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.EXTRACT_DIMENSIONAL_TOPICS, Status.IN_PROGRESS) + logger.info("Starting dimensional topics extraction (L0)...") + + # Generate L0 - Call document_service to analyze all documents + logger.info("Generating L0 data...") + analyzed_docs = document_service.analyze_all_documents() + logger.info(f"Successfully analyzed {len(analyzed_docs)} documents for L0") + + # Mark step as completed + self.progress.mark_step_status(ProcessStep.EXTRACT_DIMENSIONAL_TOPICS, Status.COMPLETED) + logger.info("Dimensional topics extraction (L0) completed successfully") + return True + + except Exception as e: + logger.error(f"Extract dimensional topics (L0) failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.EXTRACT_DIMENSIONAL_TOPICS, Status.FAILED) + return False + + def generate_biography(self) -> bool: + """Generate biography using L1 data""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.GENERATE_BIOGRAPHY, Status.IN_PROGRESS) + logger.info("Starting biography generation...") + + # Generate L1 data and biography + logger.info("Generating L1 data and biography...") + l1_data = generate_l1_from_l0() + logger.info("Successfully generated L1 data and biography") + + # Store L1 data + with DatabaseSession.session() as session: + store_l1_data(session, l1_data) + + # Mark step as completed + self.progress.mark_step_status(ProcessStep.GENERATE_BIOGRAPHY, Status.COMPLETED) + logger.info("Biography generation completed successfully") + return True + + except Exception as e: + logger.error(f"Biography generation failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.GENERATE_BIOGRAPHY, Status.FAILED) + return False + + def model_download(self) -> bool: + """Download model""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.MODEL_DOWNLOAD, Status.IN_PROGRESS) + # Directly call save_hf_model function to download model + logger.info(f"Starting model download: {self.model_name}") + + # Start monitoring the download progress in a separate thread + monitor_thread = threading.Thread(target=self._monitor_model_download) + monitor_thread.daemon = True + monitor_thread.start() + + # Start the actual download + model_path = save_hf_model(self.model_name) + + if model_path and os.path.exists(model_path): + logger.info(f"Model downloaded successfully to {model_path}") + self.progress.mark_step_status(ProcessStep.MODEL_DOWNLOAD, Status.COMPLETED) + return True + else: + logger.error(f"Model path does not exist after download: {model_path}") + self.progress.mark_step_status(ProcessStep.MODEL_DOWNLOAD, Status.FAILED) + return False + + except Exception as e: + logger.error(f"Download model failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.MODEL_DOWNLOAD, Status.FAILED) + return False + + def map_your_entity_network(self)->bool: + """Map entity network using notes and basic info""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.MAP_ENTITY_NETWORK, Status.IN_PROGRESS) + logger.info("Starting entity network mapping...") + + # Get or prepare L2 data + self._prepare_l2_data() + + l2_generator = L2Generator( + data_path=os.path.join(os.getcwd(), "resources") + ) + l2_generator.data_preprocess(self.l2_data["notes"], self.l2_data["basic_info"]) + + self.progress.mark_step_status(ProcessStep.MAP_ENTITY_NETWORK, Status.COMPLETED) + logger.info("Entity network mapping completed successfully") + return True + + except Exception as e: + logger.error(f"Map entity network failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.MAP_ENTITY_NETWORK, Status.FAILED) + self._cleanup_resources() + return False + + def decode_preference_patterns(self)->bool: + """Decode preference patterns using notes and related data""" + try: + params_manager = TrainingParamsManager() + training_params = params_manager.get_latest_training_params() + concurrency_threads = training_params.get("concurrency_threads") + data_synthesis_mode = training_params.get("data_synthesis_mode") + os.environ["CONCURRENCY_THREADS"] = str(concurrency_threads) + os.environ["DATA_SYNTHESIS_MODE"] = data_synthesis_mode + + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.DECODE_PREFERENCE_PATTERNS, Status.IN_PROGRESS) + logger.info("Starting preference patterns decoding...") + # Get or prepare L2 data + self._prepare_l2_data() + + # Use data from l2_data dictionary + training_params = TrainingParamsManager.get_latest_training_params() + L2Generator(is_cot=training_params.get("is_cot", False)).gen_preference_data( + self.l2_data["notes"], + self.l2_data["basic_info"], + self.l2_data["data_output_base_dir"], + self.l2_data["topics_path"], + self.l2_data["entitys_path"], + self.l2_data["graph_path"], + self.l2_data["config_path"] + ) + + self.progress.mark_step_status(ProcessStep.DECODE_PREFERENCE_PATTERNS, Status.COMPLETED) + logger.info("Preference patterns decoding completed successfully") + return True + + except Exception as e: + logger.error(f"Decode preference patterns failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.DECODE_PREFERENCE_PATTERNS, Status.FAILED) + return False + + def reinforce_identity(self)->bool: + """Reinforce identity using notes and related data""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.REINFORCE_IDENTITY, Status.IN_PROGRESS) + logger.info("Starting identity reinforcement...") + # Get or prepare L2 data + self._prepare_l2_data() + + # Get training parameters + training_params = TrainingParamsManager.get_latest_training_params() + # Use data from l2_data dictionary + l2_generator = L2Generator( + data_path=os.path.join(os.getcwd(), "resources"), is_cot=training_params.get("is_cot", False) + ) + l2_generator.gen_selfqa_data( + self.l2_data["notes"], + self.l2_data["basic_info"], + self.l2_data["data_output_base_dir"], + self.l2_data["topics_path"], + self.l2_data["entitys_path"], + self.l2_data["graph_path"], + self.l2_data["config_path"] + ) + + self.progress.mark_step_status(ProcessStep.REINFORCE_IDENTITY, Status.COMPLETED) + logger.info("Identity reinforcement completed successfully") + return True + + except Exception as e: + logger.error(f"Reinforce identity failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.REINFORCE_IDENTITY, Status.FAILED) + return False + + def _cleanup_resources(self): + """Clean up resources to prevent memory leaks""" + logger.info("Cleaning up resources to prevent memory leaks") + + # Clean up large data structures in l2_data dictionary + for key in self.l2_data: + self.l2_data[key] = None + + self.l2_data_prepared = False + + # Force garbage collection + gc.collect() + + # Log memory usage after cleanup + process = psutil.Process(os.getpid()) + memory_info = process.memory_info() + logger.info(f"Memory usage after cleanup: {memory_info.rss / 1024 / 1024:.2f} MB") + + def augment_content_retention(self) -> bool: + """Augment content retention using notes, basic info and graph data""" + try: + # Mark step as in progress + self.progress.mark_step_status(ProcessStep.AUGMENT_CONTENT_RETENTION, Status.IN_PROGRESS) + logger.info("Starting content retention augmentation...") + # Get or prepare L2 data + self._prepare_l2_data() + + # Get training parameters + training_params = TrainingParamsManager.get_latest_training_params() + # Use data from l2_data dictionary + l2_generator = L2Generator(data_path=os.path.join(os.getcwd(), "resources"), is_cot=training_params.get("is_cot", False)) + l2_generator.gen_diversity_data( + self.l2_data["notes"], + self.l2_data["basic_info"], + self.l2_data["data_output_base_dir"], + self.l2_data["topics_path"], + self.l2_data["entitys_path"], + self.l2_data["graph_path"], + self.l2_data["config_path"] + ) + l2_generator.merge_json_files(self.l2_data["data_output_base_dir"]) + + self.progress.mark_step_status(ProcessStep.AUGMENT_CONTENT_RETENTION, Status.COMPLETED) + logger.info("Content retention augmentation completed successfully") + return True + + except Exception as e: + logger.error(f"Augment content retention failed: {str(e)}") + self.progress.mark_step_status(ProcessStep.AUGMENT_CONTENT_RETENTION, Status.FAILED) + return False + + def _prepare_l2_data(self): + """Prepare L2 data for training""" + if self.l2_data_prepared: + logger.info("L2 data already prepared, skipping preparation") + return + + try: + logger.info("Preparing L2 data...") + + # Get notes from database + note_service = NoteService() + notes = note_service.get_all_notes() + + # Get basic info + status_bio = get_latest_status_bio() + global_bio = get_latest_global_bio() + + # Combine into basic info dictionary + basic_info = { + "status_bio": status_bio, + "global_bio": global_bio + } + + # Create output directory + data_output_base_dir = os.path.join(os.getcwd(), "data", "l2_data") + os.makedirs(data_output_base_dir, exist_ok=True) + + # Set paths for L2 data + topics_path = os.path.join(data_output_base_dir, "topics.json") + entitys_path = os.path.join(data_output_base_dir, "entitys.json") + graph_path = os.path.join(data_output_base_dir, "graph.json") + config_path = os.path.join(data_output_base_dir, "config.json") + + # Store in l2_data dictionary + self.l2_data["notes"] = notes + self.l2_data["basic_info"] = basic_info + self.l2_data["data_output_base_dir"] = data_output_base_dir + self.l2_data["topics_path"] = topics_path + self.l2_data["entitys_path"] = entitys_path + self.l2_data["graph_path"] = graph_path + self.l2_data["config_path"] = config_path + + self.l2_data_prepared = True + logger.info("L2 data preparation completed successfully") + + except Exception as e: + logger.error(f"L2 data preparation failed: {str(e)}") + self._cleanup_resources() + raise + + def _monitor_model_download(self): + """Monitor model download progress""" + try: + logger.info("Starting model download monitoring...") + + # Check every 5 seconds + while True: + # Sleep first to give download time to start + time.sleep(5) + + # Check if download is still in progress + if self.progress.get_step_status(ProcessStep.MODEL_DOWNLOAD) != Status.IN_PROGRESS: + logger.info("Model download monitoring stopped - download completed or failed") + break + + # Log current memory usage + process = psutil.Process(os.getpid()) + memory_info = process.memory_info() + logger.info(f"Current memory usage during download: {memory_info.rss / 1024 / 1024:.2f} MB") + + except Exception as e: + logger.error(f"Error in model download monitoring: {str(e)}") + + def check_training_condition(self) -> bool: + """Check if training conditions are met""" + try: + # Check if documents exist + documents = self.list_documents() + if not documents: + logger.error("No documents found, training cannot proceed") + return False + + # Check if model name is valid + if not self.model_name: + logger.error("Model name is not set, training cannot proceed") + return False + + return True + except Exception as e: + logger.error(f"Error checking training conditions: {str(e)}") + return False + + def reset_progress(self): + """Reset training progress""" + try: + logger.info("Resetting training progress...") + self.progress.reset_progress() + logger.info("Training progress reset successfully") + except Exception as e: + logger.error(f"Error resetting training progress: {str(e)}") + + def start_process(self): + """Start the training process""" + try: + logger.info("Starting training process...") + + # Reset stop flag + self.is_stopped = False + + # Reset progress + self.reset_progress() + + # Check training conditions + if not self.check_training_condition(): + logger.error("Training conditions not met, aborting process") + return False + + # Start the process + self.progress.mark_status(Status.IN_PROGRESS) + + # Execute each step in sequence + steps = [ + ("List documents", self.list_documents), + ("Generate document embeddings", self.generate_document_embeddings), + ("Process chunks", self.process_chunks), + ("Generate chunk embeddings", self.chunk_embedding), + ("Extract dimensional topics", self.extract_dimensional_topics), + ("Generate biography", self.generate_biography), + ("Download model", self.model_download), + ("Map entity network", self.map_your_entity_network), + ("Decode preference patterns", self.decode_preference_patterns), + ("Reinforce identity", self.reinforce_identity), + ("Augment content retention", self.augment_content_retention), + ] + + for step_name, step_func in steps: + if self.is_stopped: + logger.info("Process stopped by user") + self.progress.mark_status(Status.STOPPED) + return False + + logger.info(f"Executing step: {step_name}") + self.current_step = step_name + + # Execute step + success = step_func() + + if not success: + logger.error(f"Step '{step_name}' failed, stopping process") + self.progress.mark_status(Status.FAILED) + return False + + # All steps completed successfully + logger.info("All steps completed successfully") + self.progress.mark_status(Status.COMPLETED) + return True + + except Exception as e: + logger.error(f"Error in training process: {str(e)}") + self.progress.mark_status(Status.FAILED) + return False + + def stop_process(self): + """Stop the training process""" + try: + logger.info("Stopping training process...") + self.is_stopped = True + logger.info("Training process stop flag set") + return True + except Exception as e: + logger.error(f"Error stopping training process: {str(e)}") + return False \ No newline at end of file