import os import json import hashlib from threading import Lock from typing import List, Dict, Any, Optional class ModelPathManager: """ 根据定义的规则管理模型训练产物的保存路径。 此类旨在集中处理所有与文件系统交互的路径生成逻辑， 确保整个应用程序遵循统一的模型保存标准。 """ def __init__(self, base_dir: str = 'saved_models'): """ 初始化路径管理器。 Args: base_dir (str): 所有模型保存的根目录。 """ # 始终使用相对于项目根目录的相对路径 self.base_dir = base_dir self.versions_file = os.path.join(self.base_dir, 'versions.json') self.lock = Lock() # 确保根目录存在 os.makedirs(self.base_dir, exist_ok=True) def _hash_ids(self, ids: List[str]) -> str: """ 对ID列表进行排序和哈希，生成一个稳定的、简短的哈希值。 Args: ids (List[str]): 需要哈希的ID列表。 Returns: str: 代表该ID集合的10位短哈希字符串。 """ if not ids: return 'none' # 排序以确保对于相同集合的ID，即使顺序不同，结果也一样 sorted_ids = sorted([str(i) for i in ids]) id_string = ",".join(sorted_ids) # 使用SHA256生成哈希值并截取前10位 return hashlib.sha256(id_string.encode('utf-8')).hexdigest()[:10] def _generate_identifier(self, training_mode: str, **kwargs: Any) -> str: """ 根据训练模式和参数生成模型的唯一标识符 (identifier)。 这个标识符将作为版本文件中的key，并用于构建目录路径。 Args: training_mode (str): 训练模式 ('product', 'store', 'global')。 **kwargs: 从API请求中传递的参数字典。 Returns: str: 模型的唯一标识符。 Raises: ValueError: 如果缺少必要的参数。 """ if training_mode == 'product': product_id = kwargs.get('product_id') if not product_id: raise ValueError("按药品训练模式需要 'product_id'。") # 对于药品训练，数据范围由 store_id 定义 store_id = kwargs.get('store_id') scope = store_id if store_id is not None else 'all' return f"product_{product_id}_scope_{scope}" elif training_mode == 'store': store_id = kwargs.get('store_id') if not store_id: raise ValueError("按店铺训练模式需要 'store_id'。") product_scope = kwargs.get('product_scope', 'all') if product_scope == 'specific': product_ids = kwargs.get('product_ids') if not product_ids: raise ValueError("店铺训练选择 specific 范围时需要 'product_ids'。") # 如果只有一个ID，直接使用ID；否则使用哈希 scope = product_ids[0] if len(product_ids) == 1 else self._hash_ids(product_ids) else: scope = 'all' return f"store_{store_id}_products_{scope}" elif training_mode == 'global': training_scope = kwargs.get('training_scope', 'all') if training_scope in ['all', 'all_stores_all_products']: scope_part = 'all' elif training_scope == 'selected_stores': store_ids = kwargs.get('store_ids') if not store_ids: raise ValueError("全局训练选择 selected_stores 范围时需要 'store_ids'。") # 如果只有一个ID，直接使用ID；否则使用哈希 scope_id = store_ids[0] if len(store_ids) == 1 else self._hash_ids(store_ids) scope_part = f"stores_{scope_id}" elif training_scope == 'selected_products': product_ids = kwargs.get('product_ids') if not product_ids: raise ValueError("全局训练选择 selected_products 范围时需要 'product_ids'。") # 如果只有一个ID，直接使用ID；否则使用哈希 scope_id = product_ids[0] if len(product_ids) == 1 else self._hash_ids(product_ids) scope_part = f"products_{scope_id}" elif training_scope == 'custom': store_ids = kwargs.get('store_ids') product_ids = kwargs.get('product_ids') if not store_ids or not product_ids: raise ValueError("全局训练选择 custom 范围时需要 'store_ids' 和 'product_ids'。") s_id = store_ids[0] if len(store_ids) == 1 else self._hash_ids(store_ids) p_id = product_ids[0] if len(product_ids) == 1 else self._hash_ids(product_ids) scope_part = f"custom_s_{s_id}_p_{p_id}" else: raise ValueError(f"未知的全局训练范围: {training_scope}") aggregation_method = kwargs.get('aggregation_method', 'sum') return f"global_{scope_part}_{aggregation_method}" else: raise ValueError(f"未知的训练模式: {training_mode}") def get_next_version(self, identifier: str) -> int: """ 获取指定标识符的下一个版本号。 此方法是线程安全的。 Args: identifier (str): 模型的唯一标识符。 Returns: int: 下一个可用的版本号 (从1开始)。 """ with self.lock: try: if os.path.exists(self.versions_file): with open(self.versions_file, 'r', encoding='utf-8') as f: versions_data = json.load(f) else: versions_data = {} # 如果标识符不存在，当前版本为0，下一个版本即为1 current_version = versions_data.get(identifier, 0) return current_version + 1 except (IOError, json.JSONDecodeError) as e: # 如果文件损坏或读取失败，从0开始 print(f"警告: 读取版本文件 '{self.versions_file}' 失败: {e}。将从版本1开始。") return 1 def save_version_info(self, identifier: str, new_version: int): """ 训练成功后，更新版本文件。 此方法是线程安全的。 Args: identifier (str): 模型的唯一标识符。 new_version (int): 要保存的新的版本号。 """ with self.lock: try: if os.path.exists(self.versions_file): with open(self.versions_file, 'r', encoding='utf-8') as f: versions_data = json.load(f) else: versions_data = {} versions_data[identifier] = new_version with open(self.versions_file, 'w', encoding='utf-8') as f: json.dump(versions_data, f, indent=4, ensure_ascii=False) except (IOError, json.JSONDecodeError) as e: print(f"错误: 保存版本信息到 '{self.versions_file}' 失败: {e}") # 在这种情况下，可以选择抛出异常或采取其他恢复措施 raise def get_model_paths(self, training_mode: str, model_type: str, **kwargs: Any) -> Dict[str, Any]: """ 主入口函数：为一次新的训练获取所有相关路径和版本信息。 此方法遵循扁平化文件存储规范，将逻辑路径编码到文件名中。 Args: training_mode (str): 训练模式 ('product', 'store', 'global')。 model_type (str): 模型类型 (e.g., 'mlstm', 'kan')。 **kwargs: 从API请求中传递的参数字典。 Returns: Dict[str, Any]: 一个包含所有路径和关键信息的字典。 """ # 1. 生成不含模型类型和版本的核心标识符，并将其中的分隔符替换为下划线 # 例如：product/P001/all -> product_P001_all base_identifier = self._generate_identifier(training_mode, **kwargs) # 规范化处理，将 'scope' 'products' 等关键字替换为更简洁的形式 # 例如 product_P001_scope_all -> product_P001_all core_prefix = base_identifier.replace('_scope_', '_').replace('_products_', '_') # 2. 构建用于版本控制的完整标识符 (不含版本号) # 例如: product_P001_all_mlstm version_control_identifier = f"{core_prefix}_{model_type}" # 3. 获取下一个版本号 next_version = self.get_next_version(version_control_identifier) version_str = f"v{next_version}" # 4. 构建最终的文件名前缀，包含版本号 # 例如: product_P001_all_mlstm_v2 filename_prefix = f"{version_control_identifier}_{version_str}" # 5. 确保 `saved_models` 和 `saved_models/checkpoints` 目录存在 checkpoints_base_dir = os.path.join(self.base_dir, 'checkpoints') os.makedirs(self.base_dir, exist_ok=True) os.makedirs(checkpoints_base_dir, exist_ok=True) # 6. 构建并返回包含所有扁平化路径和关键信息的字典 return { "identifier": version_control_identifier, # 用于版本控制的key "filename_prefix": filename_prefix, # 用于数据库和文件查找 "version": next_version, "base_dir": self.base_dir, "model_path": os.path.join(self.base_dir, f"{filename_prefix}_model.pth"), "metadata_path": os.path.join(self.base_dir, f"{filename_prefix}_metadata.json"), "loss_curve_path": os.path.join(self.base_dir, f"{filename_prefix}_loss_curve.png"), "checkpoint_dir": checkpoints_base_dir, # 指向公共的检查点目录 "best_checkpoint_path": os.path.join(checkpoints_base_dir, f"{filename_prefix}_checkpoint_best.pth"), # 为动态epoch检查点提供一个格式化模板 "epoch_checkpoint_template": os.path.join(checkpoints_base_dir, f"{filename_prefix}_checkpoint_epoch_{{N}}.pth") } def get_model_path_for_prediction(self, training_mode: str, model_type: str, version: int, **kwargs: Any) -> Optional[str]: """ 获取用于预测的已存在模型的完整路径 (遵循扁平化规范)。 Args: training_mode (str): 训练模式。 model_type (str): 模型类型。 version (int): 模型版本号。 **kwargs: 其他用于定位模型的参数。 Returns: Optional[str]: 模型的完整路径，如果不存在则返回None。 """ # 1. 生成不含模型类型和版本的核心标识符 base_identifier = self._generate_identifier(training_mode, **kwargs) core_prefix = base_identifier.replace('_scope_', '_').replace('_products_', '_') # 2. 构建用于版本控制的标识符 version_control_identifier = f"{core_prefix}_{model_type}" # 3. 构建完整的文件名前缀 version_str = f"v{version}" filename_prefix = f"{version_control_identifier}_{version_str}" # 4. 构建模型文件的完整路径 model_path = os.path.join(self.base_dir, f"{filename_prefix}_model.pth") return model_path if os.path.exists(model_path) else None