""" 统一模型管理工具 处理模型文件的统一命名、存储和检索 遵循层级式目录结构和文件版本管理规则 """ import os import json import torch import glob from datetime import datetime from typing import List, Dict, Optional, Any from threading import Lock from core.config import DEFAULT_MODEL_DIR class ModelManager: """ 统一模型管理器，采用结构化目录和版本文件进行管理。 """ VERSION_FILE = 'versions.json' def __init__(self, model_dir: str = DEFAULT_MODEL_DIR): self.model_dir = os.path.abspath(model_dir) self.versions_path = os.path.join(self.model_dir, self.VERSION_FILE) self._lock = Lock() self.ensure_model_dir() def ensure_model_dir(self): """确保模型根目录存在""" os.makedirs(self.model_dir, exist_ok=True) def _read_versions(self) -> Dict[str, int]: """线程安全地读取版本文件""" with self._lock: if not os.path.exists(self.versions_path): return {} try: with open(self.versions_path, 'r', encoding='utf-8') as f: return json.load(f) except (json.JSONDecodeError, IOError): return {} def _write_versions(self, versions: Dict[str, int]): """线程安全地写入版本文件""" with self._lock: with open(self.versions_path, 'w', encoding='utf-8') as f: json.dump(versions, f, indent=2) def get_model_identifier(self, model_type: str, training_mode: str, scope: str, aggregation_method: Optional[str] = None) -> str: """ 生成模型的唯一标识符，用于版本文件中的key。 """ if training_mode == 'global': return f"{training_mode}_{scope}_{aggregation_method}_{model_type}" return f"{training_mode}_{scope}_{model_type}" def get_next_version_number(self, model_identifier: str) -> int: """ 获取指定模型的下一个版本号（整数）。 """ versions = self._read_versions() current_version = versions.get(model_identifier, 0) return current_version + 1 def update_version(self, model_identifier: str, new_version: int): """ 更新模型的最新版本号。 """ versions = self._read_versions() versions[model_identifier] = new_version self._write_versions(versions) def get_model_version_path(self, model_type: str, training_mode: str, scope: str, version: int, aggregation_method: Optional[str] = None) -> str: """ 根据新规则生成模型版本目录的完整路径。 """ base_path = os.path.join(self.model_dir, training_mode, scope) if training_mode == 'global' and aggregation_method: base_path = os.path.join(base_path, str(aggregation_method)) version_path = os.path.join(base_path, model_type, f'v{version}') return version_path def save_model_artifact(self, artifact_data: Any, artifact_name: str, model_version_path: str): """ 在指定的模型版本目录下保存一个产物。 Args: artifact_data: 要保存的数据 (e.g., model state dict, figure object). artifact_name: 标准化的产物文件名 (e.g., 'model.pth', 'loss_curve.png'). model_version_path: 模型版本目录的路径. """ os.makedirs(model_version_path, exist_ok=True) full_path = os.path.join(model_version_path, artifact_name) if artifact_name.endswith('.pth'): torch.save(artifact_data, full_path) elif artifact_name.endswith('.png') and hasattr(artifact_data, 'savefig'): artifact_data.savefig(full_path, dpi=300, bbox_inches='tight') elif artifact_name.endswith('.json'): with open(full_path, 'w', encoding='utf-8') as f: json.dump(artifact_data, f, indent=2, ensure_ascii=False) else: raise ValueError(f"不支持的产物类型: {artifact_name}") print(f"产物已保存: {full_path}") def list_models(self, page: Optional[int] = None, page_size: Optional[int] = None) -> Dict: """ 通过扫描目录结构来列出所有模型。 """ all_models = [] for training_mode in os.listdir(self.model_dir): mode_path = os.path.join(self.model_dir, training_mode) if not os.path.isdir(mode_path) or training_mode == 'checkpoints' or training_mode == self.VERSION_FILE: continue for scope in os.listdir(mode_path): scope_path = os.path.join(mode_path, scope) if not os.path.isdir(scope_path): continue is_global_agg_level = False if training_mode == 'global' and os.listdir(scope_path): try: first_item_path = os.path.join(scope_path, os.listdir(scope_path)[0]) if os.path.isdir(first_item_path): is_global_agg_level = True except IndexError: continue if is_global_agg_level: for agg_method in os.listdir(scope_path): agg_path = os.path.join(scope_path, agg_method) if not os.path.isdir(agg_path): continue for model_type in os.listdir(agg_path): type_path = os.path.join(agg_path, model_type) if not os.path.isdir(type_path): continue for version_folder in os.listdir(type_path): if version_folder.startswith('v'): version_path = os.path.join(type_path, version_folder) model_info = self._parse_info_from_path(version_path) if model_info: all_models.append(model_info) else: for model_type in os.listdir(scope_path): type_path = os.path.join(scope_path, model_type) if not os.path.isdir(type_path): continue for version_folder in os.listdir(type_path): if version_folder.startswith('v'): version_path = os.path.join(type_path, version_folder) model_info = self._parse_info_from_path(version_path) if model_info: all_models.append(model_info) total_count = len(all_models) if page and page_size: start_index = (page - 1) * page_size end_index = start_index + page_size paginated_models = all_models[start_index:end_index] else: paginated_models = all_models return { 'models': paginated_models, 'pagination': { 'total': total_count, 'page': page or 1, 'page_size': page_size or total_count, 'total_pages': (total_count + page_size - 1) // page_size if page_size and page_size > 0 else 1, } } def _parse_info_from_path(self, version_path: str) -> Optional[Dict]: """从版本目录路径解析模型信息""" try: norm_path = os.path.normpath(version_path) norm_model_dir = os.path.normpath(self.model_dir) relative_path = os.path.relpath(norm_path, norm_model_dir) parts = relative_path.split(os.sep) info = { 'model_path': version_path, 'version': parts[-1], 'model_type': parts[-2] } training_mode = parts[0] info['training_mode'] = training_mode if training_mode == 'global': info['scope'] = parts[1] info['aggregation_method'] = parts[2] info['model_identifier'] = self.get_model_identifier(info['model_type'], training_mode, info['scope'], info['aggregation_method']) else: info['scope'] = parts[1] info['aggregation_method'] = None info['model_identifier'] = self.get_model_identifier(info['model_type'], training_mode, info['scope']) metadata_path = os.path.join(version_path, 'metadata.json') if os.path.exists(metadata_path): with open(metadata_path, 'r', encoding='utf-8') as f: metadata = json.load(f) info.update(metadata) return info except (IndexError, IOError) as e: print(f"解析路径失败 {version_path}: {e}") return None # 全局模型管理器实例 model_manager = ModelManager()