|
@@ -1,11 +1,15 @@
|
|
-from pymongo import MongoClient, UpdateOne, DESCENDING
|
|
|
|
|
|
+import pymongo
|
|
|
|
+from pymongo import MongoClient, UpdateOne, DESCENDING, ASCENDING
|
|
|
|
+from pymongo.errors import PyMongoError
|
|
import pandas as pd
|
|
import pandas as pd
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy import create_engine
|
|
import pickle
|
|
import pickle
|
|
from io import BytesIO
|
|
from io import BytesIO
|
|
import joblib
|
|
import joblib
|
|
-import h5py
|
|
|
|
|
|
+import h5py, os, io
|
|
import tensorflow as tf
|
|
import tensorflow as tf
|
|
|
|
+from typing import Dict, Any, Optional, Union, Tuple
|
|
|
|
+import tempfile
|
|
|
|
|
|
def get_data_from_mongo(args):
|
|
def get_data_from_mongo(args):
|
|
mongodb_connection = "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/"
|
|
mongodb_connection = "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/"
|
|
@@ -141,133 +145,348 @@ def insert_pickle_model_into_mongo(model, args):
|
|
print("model inserted successfully!")
|
|
print("model inserted successfully!")
|
|
|
|
|
|
|
|
|
|
-def insert_h5_model_into_mongo(model,feature_scaler_bytes,target_scaler_bytes ,args):
|
|
|
|
- mongodb_connection,mongodb_database,scaler_table,model_table,model_name = ("mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",
|
|
|
|
- args['mongodb_database'],args['scaler_table'],args['model_table'],args['model_name'])
|
|
|
|
- client = MongoClient(mongodb_connection)
|
|
|
|
- db = client[mongodb_database]
|
|
|
|
- if scaler_table in db.list_collection_names():
|
|
|
|
- db[scaler_table].drop()
|
|
|
|
- print(f"Collection '{scaler_table} already exist, deleted successfully!")
|
|
|
|
- collection = db[scaler_table] # 集合名称
|
|
|
|
- # Save the scalers in MongoDB as binary data
|
|
|
|
- collection.insert_one({
|
|
|
|
- "feature_scaler": feature_scaler_bytes.read(),
|
|
|
|
- "target_scaler": target_scaler_bytes.read()
|
|
|
|
- })
|
|
|
|
- print("scaler_model inserted successfully!")
|
|
|
|
- if model_table in db.list_collection_names():
|
|
|
|
- db[model_table].drop()
|
|
|
|
- print(f"Collection '{model_table} already exist, deleted successfully!")
|
|
|
|
- model_table = db[model_table]
|
|
|
|
- # 创建 BytesIO 缓冲区
|
|
|
|
- model_buffer = BytesIO()
|
|
|
|
- # 将模型保存为 HDF5 格式到内存 (BytesIO)
|
|
|
|
- model.save(model_buffer, save_format='h5')
|
|
|
|
- # 将指针移到缓冲区的起始位置
|
|
|
|
- model_buffer.seek(0)
|
|
|
|
- # 获取模型的二进制数据
|
|
|
|
- model_data = model_buffer.read()
|
|
|
|
- # 将模型保存到 MongoDB
|
|
|
|
- model_table.insert_one({
|
|
|
|
- "model_name": model_name,
|
|
|
|
- "model_data": model_data
|
|
|
|
- })
|
|
|
|
- print("模型成功保存到 MongoDB!")
|
|
|
|
-
|
|
|
|
-def insert_trained_model_into_mongo(model, args):
|
|
|
|
- mongodb_connection,mongodb_database,model_table,model_name = ("mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",
|
|
|
|
- args['mongodb_database'],args['model_table'],args['model_name'])
|
|
|
|
-
|
|
|
|
- gen_time, params_json, descr = args['gen_time'], args['params'], args['descr']
|
|
|
|
- client = MongoClient(mongodb_connection)
|
|
|
|
- db = client[mongodb_database]
|
|
|
|
- if model_table in db.list_collection_names():
|
|
|
|
- db[model_table].drop()
|
|
|
|
- print(f"Collection '{model_table} already exist, deleted successfully!")
|
|
|
|
- model_table = db[model_table]
|
|
|
|
- # 创建 BytesIO 缓冲区
|
|
|
|
- model_buffer = BytesIO()
|
|
|
|
- # 将模型保存为 HDF5 格式到内存 (BytesIO)
|
|
|
|
- model.save(model_buffer, save_format='h5')
|
|
|
|
- # 将指针移到缓冲区的起始位置
|
|
|
|
- model_buffer.seek(0)
|
|
|
|
- # 获取模型的二进制数据
|
|
|
|
- model_data = model_buffer.read()
|
|
|
|
- # 将模型保存到 MongoDB
|
|
|
|
- model_table.insert_one({
|
|
|
|
- "model_name": model_name,
|
|
|
|
- "model_data": model_data,
|
|
|
|
- "gen_time": gen_time,
|
|
|
|
- "params": params_json,
|
|
|
|
- "descr": descr
|
|
|
|
- })
|
|
|
|
- print("模型成功保存到 MongoDB!")
|
|
|
|
-
|
|
|
|
-def insert_scaler_model_into_mongo(feature_scaler_bytes, scaled_target_bytes, args):
|
|
|
|
- mongodb_connection,mongodb_database,scaler_table,model_table,model_name = ("mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",
|
|
|
|
- args['mongodb_database'],args['scaler_table'],args['model_table'],args['model_name'])
|
|
|
|
- gen_time = args['gen_time']
|
|
|
|
- client = MongoClient(mongodb_connection)
|
|
|
|
- db = client[mongodb_database]
|
|
|
|
- if scaler_table in db.list_collection_names():
|
|
|
|
- db[scaler_table].drop()
|
|
|
|
- print(f"Collection '{scaler_table} already exist, deleted successfully!")
|
|
|
|
- collection = db[scaler_table] # 集合名称
|
|
|
|
- # Save the scalers in MongoDB as binary data
|
|
|
|
- collection.insert_one({
|
|
|
|
- "model_name": model_name,
|
|
|
|
- "gent_time": gen_time,
|
|
|
|
- "feature_scaler": feature_scaler_bytes.read(),
|
|
|
|
- "target_scaler": scaled_target_bytes.read()
|
|
|
|
- })
|
|
|
|
- print("scaler_model inserted successfully!")
|
|
|
|
-
|
|
|
|
-
|
|
|
|
-def get_h5_model_from_mongo(args, custom=None):
|
|
|
|
- mongodb_connection,mongodb_database,model_table,model_name = "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/",args['mongodb_database'],args['model_table'],args['model_name']
|
|
|
|
- client = MongoClient(mongodb_connection)
|
|
|
|
- # 选择数据库(如果数据库不存在,MongoDB 会自动创建)
|
|
|
|
- db = client[mongodb_database]
|
|
|
|
- collection = db[model_table] # 集合名称
|
|
|
|
-
|
|
|
|
- # 查询 MongoDB 获取模型数据
|
|
|
|
- model_doc = collection.find_one({"model_name": model_name}, sort=[('gen_time', DESCENDING)])
|
|
|
|
- if model_doc:
|
|
|
|
- model_data = model_doc['model_data'] # 获取模型的二进制数据
|
|
|
|
- # 将二进制数据加载到 BytesIO 缓冲区
|
|
|
|
- model_buffer = BytesIO(model_data)
|
|
|
|
- # 从缓冲区加载模型
|
|
|
|
- # 使用 h5py 和 BytesIO 从内存中加载模型
|
|
|
|
- with h5py.File(model_buffer, 'r') as f:
|
|
|
|
- model = tf.keras.models.load_model(f, custom_objects=custom)
|
|
|
|
- print(f"{model_name}模型成功从 MongoDB 加载!")
|
|
|
|
- client.close()
|
|
|
|
- return model
|
|
|
|
- else:
|
|
|
|
- print(f"未找到model_name为 {model_name} 的模型。")
|
|
|
|
- client.close()
|
|
|
|
- return None
|
|
|
|
|
|
+def insert_trained_model_into_mongo(model: tf.keras.Model, args: Dict[str, Any]) -> str:
|
|
|
|
+ """
|
|
|
|
+ 将训练好的H5模型插入MongoDB,自动维护集合容量不超过50个模型
|
|
|
|
+ 参数:
|
|
|
|
+ model : keras模型 - 训练好的Keras模型
|
|
|
|
+ args : dict - 包含以下键的字典:
|
|
|
|
+ mongodb_database: 数据库名称
|
|
|
|
+ model_table: 集合名称
|
|
|
|
+ model_name: 模型名称
|
|
|
|
+ gen_time: 模型生成时间(datetime对象)
|
|
|
|
+ params: 模型参数(JSON可序列化对象)
|
|
|
|
+ descr: 模型描述文本
|
|
|
|
+ """
|
|
|
|
+ # ------------------------- 参数校验 -------------------------
|
|
|
|
+ required_keys = {'mongodb_database', 'model_table', 'model_name',
|
|
|
|
+ 'gen_time', 'params', 'descr'}
|
|
|
|
+ if missing := required_keys - args.keys():
|
|
|
|
+ raise ValueError(f"缺少必要参数: {missing}")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 配置解耦 -------------------------
|
|
|
|
+ # 从环境变量获取连接信息(更安全)
|
|
|
|
+ mongodb_connection = os.getenv("MONGO_URI", "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 资源初始化 -------------------------
|
|
|
|
+ fd, temp_path = None, None
|
|
|
|
+ client = None
|
|
|
|
+
|
|
|
|
+ try:
|
|
|
|
+ # ------------------------- 临时文件处理 -------------------------
|
|
|
|
+ fd, temp_path = tempfile.mkstemp(suffix='.h5')
|
|
|
|
+ os.close(fd) # 立即释放文件锁
|
|
|
|
+
|
|
|
|
+ # ------------------------- 模型保存 -------------------------
|
|
|
|
+ try:
|
|
|
|
+ model.save(temp_path, save_format='h5')
|
|
|
|
+ except Exception as e:
|
|
|
|
+ raise RuntimeError(f"模型保存失败: {str(e)}") from e
|
|
|
|
+
|
|
|
|
+ # ------------------------- 数据库连接 -------------------------
|
|
|
|
+ client = MongoClient(mongodb_connection)
|
|
|
|
+ db = client[args['mongodb_database']]
|
|
|
|
+ collection = db[args['model_table']]
|
|
|
|
+
|
|
|
|
+ # ------------------------- 索引检查 -------------------------
|
|
|
|
+ if "gen_time_1" not in collection.index_information():
|
|
|
|
+ collection.create_index([("gen_time", ASCENDING)], name="gen_time_1")
|
|
|
|
+ print("已创建时间索引")
|
|
|
|
|
|
|
|
+ # ------------------------- 容量控制 -------------------------
|
|
|
|
+ # 使用更高效的计数方式
|
|
|
|
+ if collection.estimated_document_count() >= 50:
|
|
|
|
+ # 原子性删除操作
|
|
|
|
+ if deleted := collection.find_one_and_delete(
|
|
|
|
+ sort=[("gen_time", ASCENDING)],
|
|
|
|
+ projection={"_id": 1, "model_name": 1, "gen_time": 1}
|
|
|
|
+ ):
|
|
|
|
+ print(f"已淘汰模型 [{deleted['model_name']}] 生成时间: {deleted['gen_time']}")
|
|
|
|
|
|
-def get_scaler_model_from_mongo(args, only_feature_scaler=False):
|
|
|
|
|
|
+ # ------------------------- 数据插入 -------------------------
|
|
|
|
+ with open(temp_path, 'rb') as f:
|
|
|
|
+ result = collection.insert_one({
|
|
|
|
+ "model_name": args['model_name'],
|
|
|
|
+ "model_data": f.read(),
|
|
|
|
+ "gen_time": args['gen_time'],
|
|
|
|
+ "params": args['params'],
|
|
|
|
+ "descr": args['descr']
|
|
|
|
+ })
|
|
|
|
+
|
|
|
|
+ print(f"✅ 模型 {args['model_name']} 保存成功 | 文档ID: {result.inserted_id}")
|
|
|
|
+ return str(result.inserted_id)
|
|
|
|
+
|
|
|
|
+ except Exception as e:
|
|
|
|
+ # ------------------------- 异常分类处理 -------------------------
|
|
|
|
+ error_type = "数据库操作" if isinstance(e, (pymongo.errors.PyMongoError, RuntimeError)) else "系统错误"
|
|
|
|
+ print(f"❌ {error_type} - 详细错误: {str(e)}")
|
|
|
|
+ raise # 根据业务需求决定是否重新抛出
|
|
|
|
+
|
|
|
|
+ finally:
|
|
|
|
+ # ------------------------- 资源清理 -------------------------
|
|
|
|
+ if client:
|
|
|
|
+ client.close()
|
|
|
|
+ if temp_path and os.path.exists(temp_path):
|
|
|
|
+ try:
|
|
|
|
+ os.remove(temp_path)
|
|
|
|
+ except PermissionError:
|
|
|
|
+ print(f"⚠️ 临时文件清理失败: {temp_path}")
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+def insert_scaler_model_into_mongo(feature_scaler_bytes: BytesIO, target_scaler_bytes: BytesIO, args: Dict[str, Any]) -> str:
|
|
"""
|
|
"""
|
|
- 根据模 型名称版本 和 生成时间 获取模型
|
|
|
|
|
|
+ 将特征缩放器存储到MongoDB,自动维护集合容量不超过50个文档
|
|
|
|
+
|
|
|
|
+ 参数:
|
|
|
|
+ feature_scaler_bytes: BytesIO - 特征缩放器字节流
|
|
|
|
+ scaled_target_bytes: BytesIO - 目标缩放器字节流
|
|
|
|
+ args : dict - 包含以下键的字典:
|
|
|
|
+ mongodb_database: 数据库名称
|
|
|
|
+ scaler_table: 集合名称
|
|
|
|
+ model_name: 关联模型名称
|
|
|
|
+ gen_time: 生成时间(datetime对象)
|
|
"""
|
|
"""
|
|
- mongodb_connection, mongodb_database, scaler_table, = ("mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/", args['mongodb_database'], args['scaler_table'])
|
|
|
|
- model_name, gen_time = args['model_name'], args['gent_time']
|
|
|
|
- client = MongoClient(mongodb_connection)
|
|
|
|
- # 选择数据库(如果数据库不存在,MongoDB 会自动创建)
|
|
|
|
- db = client[mongodb_database]
|
|
|
|
- collection = db[scaler_table] # 集合名称
|
|
|
|
- # Retrieve the scalers from MongoDB
|
|
|
|
- scaler_doc = collection.find_one({"model_name": model_name, "gen_time": gen_time})
|
|
|
|
- # Deserialize the scalers
|
|
|
|
-
|
|
|
|
- feature_scaler_bytes = BytesIO(scaler_doc["feature_scaler"])
|
|
|
|
- feature_scaler = joblib.load(feature_scaler_bytes)
|
|
|
|
- if only_feature_scaler:
|
|
|
|
- return feature_scaler
|
|
|
|
- target_scaler_bytes = BytesIO(scaler_doc["target_scaler"])
|
|
|
|
- target_scaler = joblib.load(target_scaler_bytes)
|
|
|
|
- return feature_scaler,target_scaler
|
|
|
|
|
|
+ # ------------------------- 参数校验 -------------------------
|
|
|
|
+ required_keys = {'mongodb_database', 'scaler_table', 'model_name', 'gen_time'}
|
|
|
|
+ if missing := required_keys - args.keys():
|
|
|
|
+ raise ValueError(f"缺少必要参数: {missing}")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 配置解耦 -------------------------
|
|
|
|
+ # 从环境变量获取连接信息(安全隔离凭证)
|
|
|
|
+ mongodb_conn = os.getenv("MONGO_URI", "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 输入验证 -------------------------
|
|
|
|
+ for buf, name in [(feature_scaler_bytes, "特征缩放器"),
|
|
|
|
+ (target_scaler_bytes, "目标缩放器")]:
|
|
|
|
+ if not isinstance(buf, BytesIO):
|
|
|
|
+ raise TypeError(f"{name} 必须为BytesIO类型")
|
|
|
|
+ if buf.getbuffer().nbytes == 0:
|
|
|
|
+ raise ValueError(f"{name} 字节流为空")
|
|
|
|
+
|
|
|
|
+ client = None
|
|
|
|
+ try:
|
|
|
|
+ # ------------------------- 数据库连接 -------------------------
|
|
|
|
+ client = MongoClient(mongodb_conn)
|
|
|
|
+ db = client[args['mongodb_database']]
|
|
|
|
+ collection = db[args['scaler_table']]
|
|
|
|
+
|
|
|
|
+ # ------------------------- 索引维护 -------------------------
|
|
|
|
+ if "gen_time_1" not in collection.index_information():
|
|
|
|
+ collection.create_index([("gen_time", ASCENDING)], name="gen_time_1")
|
|
|
|
+ print("⏱️ 已创建时间排序索引")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 容量控制 -------------------------
|
|
|
|
+ # 使用近似计数提升性能(误差在几十条内可接受)
|
|
|
|
+ if collection.estimated_document_count() >= 50:
|
|
|
|
+ # 原子性删除操作(保证事务完整性)
|
|
|
|
+ if deleted := collection.find_one_and_delete(
|
|
|
|
+ sort=[("gen_time", ASCENDING)],
|
|
|
|
+ projection={"_id": 1, "model_name": 1, "gen_time": 1}
|
|
|
|
+ ):
|
|
|
|
+ print(f"🗑️ 已淘汰最旧缩放器 [{deleted['model_name']}] 生成时间: {deleted['gen_time']}")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 数据插入 -------------------------
|
|
|
|
+ # 确保字节流指针位置正确
|
|
|
|
+ feature_scaler_bytes.seek(0)
|
|
|
|
+ target_scaler_bytes.seek(0)
|
|
|
|
+
|
|
|
|
+ result = collection.insert_one({
|
|
|
|
+ "model_name": args['model_name'],
|
|
|
|
+ "gen_time": args['gen_time'],
|
|
|
|
+ "feature_scaler": feature_scaler_bytes.read(),
|
|
|
|
+ "target_scaler": target_scaler_bytes.read()
|
|
|
|
+ })
|
|
|
|
+
|
|
|
|
+ print(f"✅ 缩放器 {args['model_name']} 保存成功 | 文档ID: {result.inserted_id}")
|
|
|
|
+ return str(result.inserted_id)
|
|
|
|
+
|
|
|
|
+ except Exception as e:
|
|
|
|
+ # ------------------------- 异常分类处理 -------------------------
|
|
|
|
+ error_type = "数据库操作" if isinstance(e, (pymongo.errors.PyMongoError, ValueError)) else "系统错误"
|
|
|
|
+ print(f"❌ {error_type}异常 - 详细错误: {str(e)}")
|
|
|
|
+ raise # 根据业务需求决定是否重新抛出
|
|
|
|
+
|
|
|
|
+ finally:
|
|
|
|
+ # ------------------------- 资源清理 -------------------------
|
|
|
|
+ if client:
|
|
|
|
+ client.close()
|
|
|
|
+ # 重置字节流指针(确保后续可复用)
|
|
|
|
+ feature_scaler_bytes.seek(0)
|
|
|
|
+ target_scaler_bytes.seek(0)
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+def get_h5_model_from_mongo( args: Dict[str, Any], custom_objects: Optional[Dict[str, Any]] = None) -> Optional[tf.keras.Model]:
|
|
|
|
+ """
|
|
|
|
+ 从MongoDB获取指定模型的最新版本
|
|
|
|
+
|
|
|
|
+ 参数:
|
|
|
|
+ args : dict - 包含以下键的字典:
|
|
|
|
+ mongodb_database: 数据库名称
|
|
|
|
+ model_table: 集合名称
|
|
|
|
+ model_name: 要获取的模型名称
|
|
|
|
+ custom_objects: dict - 自定义Keras对象字典
|
|
|
|
+
|
|
|
|
+ 返回:
|
|
|
|
+ tf.keras.Model - 加载成功的Keras模型
|
|
|
|
+ """
|
|
|
|
+ # ------------------------- 参数校验 -------------------------
|
|
|
|
+ required_keys = {'mongodb_database', 'model_table', 'model_name'}
|
|
|
|
+ if missing := required_keys - args.keys():
|
|
|
|
+ raise ValueError(f"❌ 缺失必要参数: {missing}")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 环境配置 -------------------------
|
|
|
|
+ mongo_uri = os.getenv("MONGO_URI", "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/")
|
|
|
|
+ client = None
|
|
|
|
+ try:
|
|
|
|
+ # ------------------------- 数据库连接 -------------------------
|
|
|
|
+ client = MongoClient(
|
|
|
|
+ mongo_uri,
|
|
|
|
+ maxPoolSize=10, # 连接池优化
|
|
|
|
+ socketTimeoutMS=5000
|
|
|
|
+ )
|
|
|
|
+ db = client[args['mongodb_database']]
|
|
|
|
+ collection = db[args['model_table']]
|
|
|
|
+
|
|
|
|
+ # ------------------------- 索引维护 -------------------------
|
|
|
|
+ index_name = "model_gen_time_idx"
|
|
|
|
+ if index_name not in collection.index_information():
|
|
|
|
+ collection.create_index(
|
|
|
|
+ [("model_name", 1), ("gen_time", DESCENDING)],
|
|
|
|
+ name=index_name
|
|
|
|
+ )
|
|
|
|
+ print("⏱️ 已创建复合索引")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 高效查询 -------------------------
|
|
|
|
+ model_doc = collection.find_one(
|
|
|
|
+ {"model_name": args['model_name']},
|
|
|
|
+ sort=[('gen_time', DESCENDING)],
|
|
|
|
+ projection={"model_data": 1, "gen_time": 1} # 获取必要字段
|
|
|
|
+ )
|
|
|
|
+
|
|
|
|
+ if not model_doc:
|
|
|
|
+ print(f"⚠️ 未找到模型 '{args['model_name']}' 的有效记录")
|
|
|
|
+ return None
|
|
|
|
+
|
|
|
|
+ # ------------------------- 内存优化加载 -------------------------
|
|
|
|
+ if model_doc:
|
|
|
|
+ model_data = model_doc['model_data'] # 获取模型的二进制数据
|
|
|
|
+ # 将二进制数据加载到 BytesIO 缓冲区
|
|
|
|
+ model_buffer = BytesIO(model_data)
|
|
|
|
+ # 从缓冲区加载模型
|
|
|
|
+ # 使用 h5py 和 BytesIO 从内存中加载模型
|
|
|
|
+ with h5py.File(model_buffer, 'r') as f:
|
|
|
|
+ model = tf.keras.models.load_model(f, custom_objects=custom_objects)
|
|
|
|
+ print(f"{args['model_name']}模型成功从 MongoDB 加载!")
|
|
|
|
+ return model
|
|
|
|
+ except tf.errors.NotFoundError as e:
|
|
|
|
+ print(f"❌ 模型结构缺失关键组件: {str(e)}")
|
|
|
|
+ raise RuntimeError("模型架构不完整") from e
|
|
|
|
+
|
|
|
|
+ except Exception as e:
|
|
|
|
+ print(f"❌ 系统异常: {str(e)}")
|
|
|
|
+ raise
|
|
|
|
+
|
|
|
|
+ finally:
|
|
|
|
+ # ------------------------- 资源清理 -------------------------
|
|
|
|
+ if client:
|
|
|
|
+ client.close()
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+def get_scaler_model_from_mongo(args: Dict[str, Any], only_feature_scaler: bool = False) -> Union[object, Tuple[object, object]]:
|
|
|
|
+ """
|
|
|
|
+ 优化版特征缩放器加载函数 - 安全高效获取最新预处理模型
|
|
|
|
+
|
|
|
|
+ 参数:
|
|
|
|
+ args : 必须包含键:
|
|
|
|
+ - mongodb_database: 数据库名称
|
|
|
|
+ - scaler_table: 集合名称
|
|
|
|
+ - model_name: 目标模型名称
|
|
|
|
+ only_feature_scaler : 是否仅返回特征缩放器
|
|
|
|
+
|
|
|
|
+ 返回:
|
|
|
|
+ 单个缩放器对象或(feature_scaler, target_scaler)元组
|
|
|
|
+
|
|
|
|
+ 异常:
|
|
|
|
+ ValueError : 参数缺失或类型错误
|
|
|
|
+ RuntimeError : 数据操作异常
|
|
|
|
+ """
|
|
|
|
+ # ------------------------- 参数校验 -------------------------
|
|
|
|
+ required_keys = {'mongodb_database', 'scaler_table', 'model_name'}
|
|
|
|
+ if missing := required_keys - args.keys():
|
|
|
|
+ raise ValueError(f"❌ 缺失必要参数: {missing}")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 环境配置 -------------------------
|
|
|
|
+ mongo_uri = os.getenv("MONGO_URI", "mongodb://root:sdhjfREWFWEF23e@192.168.1.43:30000/")
|
|
|
|
+
|
|
|
|
+ client = None
|
|
|
|
+ try:
|
|
|
|
+ # ------------------------- 数据库连接 -------------------------
|
|
|
|
+ client = MongoClient(
|
|
|
|
+ mongo_uri,
|
|
|
|
+ maxPoolSize=20, # 连接池上限
|
|
|
|
+ socketTimeoutMS=3000, # 3秒超时
|
|
|
|
+ serverSelectionTimeoutMS=5000 # 5秒服务器选择超时
|
|
|
|
+ )
|
|
|
|
+ db = client[args['mongodb_database']]
|
|
|
|
+ collection = db[args['scaler_table']]
|
|
|
|
+
|
|
|
|
+ # ------------------------- 索引维护 -------------------------
|
|
|
|
+ index_name = "model_gen_time_idx"
|
|
|
|
+ if index_name not in collection.index_information():
|
|
|
|
+ collection.create_index(
|
|
|
|
+ [("model_name", 1), ("gen_time", DESCENDING)],
|
|
|
|
+ name=index_name,
|
|
|
|
+ background=True # 后台构建避免阻塞
|
|
|
|
+ )
|
|
|
|
+ print("⏱️ 已创建特征缩放器复合索引")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 高效查询 -------------------------
|
|
|
|
+ scaler_doc = collection.find_one(
|
|
|
|
+ {"model_name": args['model_name']},
|
|
|
|
+ sort=[('gen_time', DESCENDING)],
|
|
|
|
+ projection={"feature_scaler": 1, "target_scaler": 1, "gen_time": 1}
|
|
|
|
+ )
|
|
|
|
+
|
|
|
|
+ if not scaler_doc:
|
|
|
|
+ raise RuntimeError(f"⚠️ 找不到模型 {args['model_name']} 的缩放器记录")
|
|
|
|
+
|
|
|
|
+ # ------------------------- 反序列化处理 -------------------------
|
|
|
|
+ def load_scaler(data: bytes) -> object:
|
|
|
|
+ """安全加载序列化对象"""
|
|
|
|
+ with BytesIO(data) as buffer:
|
|
|
|
+ buffer.seek(0) # 确保指针复位
|
|
|
|
+ try:
|
|
|
|
+ return joblib.load(buffer)
|
|
|
|
+ except joblib.UnpicklingError as e:
|
|
|
|
+ raise RuntimeError("反序列化失败 (可能版本不兼容)") from e
|
|
|
|
+
|
|
|
|
+ # 特征缩放器加载
|
|
|
|
+ feature_data = scaler_doc["feature_scaler"]
|
|
|
|
+ if not isinstance(feature_data, bytes):
|
|
|
|
+ raise RuntimeError("特征缩放器数据格式异常")
|
|
|
|
+ feature_scaler = load_scaler(feature_data)
|
|
|
|
+
|
|
|
|
+ if only_feature_scaler:
|
|
|
|
+ return feature_scaler
|
|
|
|
+
|
|
|
|
+ # 目标缩放器加载
|
|
|
|
+ target_data = scaler_doc["target_scaler"]
|
|
|
|
+ if not isinstance(target_data, bytes):
|
|
|
|
+ raise RuntimeError("目标缩放器数据格式异常")
|
|
|
|
+ target_scaler = load_scaler(target_data)
|
|
|
|
+
|
|
|
|
+ print(f"✅ 成功加载 {args['model_name']} 的缩放器 (版本时间: {scaler_doc.get('gen_time', '未知')})")
|
|
|
|
+ return feature_scaler, target_scaler
|
|
|
|
+
|
|
|
|
+ except PyMongoError as e:
|
|
|
|
+ raise RuntimeError(f"🔌 数据库操作失败: {str(e)}") from e
|
|
|
|
+ except RuntimeError as e:
|
|
|
|
+ raise # 直接传递已封装的异常
|
|
|
|
+ except Exception as e:
|
|
|
|
+ raise RuntimeError(f"❌ 未知系统异常: {str(e)}") from e
|
|
|
|
+ finally:
|
|
|
|
+ # ------------------------- 资源清理 -------------------------
|
|
|
|
+ if client:
|
|
|
|
+ client.close()
|
|
|
|
+
|
|
|
|
+
|