compete/app/model/task_worker.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :task_worker.py
# @Time      :2025/4/29 11:05
# @Author    :David
# @Company: shenyang JY

import logging
import pandas as pd
from app.model.tf_model_train import ModelTrainer
from app.model.tf_region_train import RegionTrainer
from app.model.material import MaterialLoader


def station_task(config):
    """场站级训练任务"""
    try:
        print("111")
        station_id = config['station_id']
        mate = MaterialLoader(base_path=config['input_file'])
        # 动态生成场站数据路径
        print("222")
        # 加载数据
        data_objects = mate.get_material(station_id)
        print("333")
        # 数据合并
        train_data = pd.merge(data_objects.nwp_v_h, data_objects.power, on=config['col_time'])
        print("444")
        # 模型训练
        # model = ModelTrainer(station_id, train_data, capacity=data_objects.cap, gpu_id=config.get('gpu_assignment'))
        model = ModelTrainer(train_data, capacity=data_objects.cap, config=config)
        model.train()
        print("555")
        return {'status': 'success', 'station_id': station_id}
    except Exception as e:
        logging.error(f"Station {station_id} failed: {str(e)}")
        return {'status': 'failed', 'station_id': station_id}


def region_task(config):
    """区域级训练任务"""
    try:
        print("111")
        station_id = config['station_id']
        mate = MaterialLoader(base_path=config['input_file'])
        # 动态生成场站数据路径
        print("222")
        # 加载数据
        data_objects = mate.get_material_region()
        print("333")
        # 数据合并
        train_data = pd.merge(data_objects.nwp_v_h, data_objects.power, on=config['col_time'])
        print("444")
        # 模型训练
        model = ModelTrainer(train_data, capacity=data_objects.cap, config=config)
        model.train()
        print("555")
        return {'status': 'success', 'station_id': station_id}
    except Exception as e:
        logging.error(f"Station {station_id} failed: {str(e)}")
        return {'status': 'failed', 'station_id': station_id}