test/models_processing/model_prediction_lightgbm/model_prediction_lightgbm.py

import lightgbm as lgb
import argparse
import pandas as pd 
import numpy as np
from pymongo import MongoClient
import pickle
from flask import Flask,request,jsonify
from waitress import serve
import time
import logging
import traceback

app = Flask('model_prediction_lightgbm——service')



def get_data_from_mongo(args):
    mongodb_connection,mongodb_database,mongodb_read_table = args['mongodb_connection'],args['mongodb_database'],args['mongodb_read_table']
    client = MongoClient(mongodb_connection)
    # 选择数据库（如果数据库不存在，MongoDB 会自动创建）
    db = client[mongodb_database]
    collection = db[mongodb_read_table]  # 集合名称
    data_from_db = collection.find()  # 这会返回一个游标（cursor）
    # 将游标转换为列表，并创建 pandas DataFrame
    df = pd.DataFrame(list(data_from_db))
    client.close()
    return df
    

def insert_data_into_mongo(res_df,args):
    mongodb_connection,mongodb_database,mongodb_write_table = args['mongodb_connection'],args['mongodb_database'],args['mongodb_write_table']
    client = MongoClient(mongodb_connection)
    db = client[mongodb_database]
    if mongodb_write_table in db.list_collection_names():
        db[mongodb_write_table].drop()
        print(f"Collection '{mongodb_write_table} already exist, deleted successfully!")
    collection = db[mongodb_write_table]  # 集合名称
    # 将 DataFrame 转为字典格式
    data_dict = res_df.to_dict("records")  # 每一行作为一个字典
    # 插入到 MongoDB
    collection.insert_many(data_dict)
    print("data inserted successfully!")
    

def model_prediction(df,args):
    mongodb_connection,mongodb_database,mongodb_model_table,model_name = args['mongodb_connection'],args['mongodb_database'],args['mongodb_model_table'],args['model_name']
    client = MongoClient(mongodb_connection)
    db = client[mongodb_database]
    collection = db[mongodb_model_table]
    model_data = collection.find_one({"model_name": model_name})
    if model_data is not None:
        model_binary = model_data['model']  # 确保这个字段是存储模型的二进制数据
        # 反序列化模型 
        model = pickle.loads(model_binary)
        df['predict'] = model.predict(df[model.feature_name()])
        print("model predict result  successfully!")
    return df


@app.route('/model_prediction_lightgbm', methods=['POST'])
def model_prediction_lightgbm():
    # 获取程序开始时间  
    start_time = time.time()  
    result = {}
    success = 0
    print("Program starts execution!")
    try:
        args = request.values.to_dict()
        print('args',args)
        logger.info(args)
        power_df = get_data_from_mongo(args)
        model = model_prediction(power_df,args)
        insert_data_into_mongo(model,args)
        success = 1
    except Exception as e:
        my_exception = traceback.format_exc()
        my_exception.replace("\n","\t")
        result['msg'] = my_exception
    end_time = time.time() 
   
    result['success'] = success
    result['args'] = args
    result['start_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))
    result['end_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))
    print("Program execution ends!")
    return result


if __name__=="__main__":  
    print("Program starts execution!")
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger = logging.getLogger("model_prediction_lightgbm log")
    from waitress import serve
    serve(app, host="0.0.0.0", port=10089)
    print("server start!")