Merge/back_rhj/app/utils/gpuwriter.py

from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine
from app.models.gpu_usage import GpuUsage
import logging

class GpuWriter:
    def __init__(self, database_url):
        """
        初始化 GpuWriter
        :param database_url: 数据库连接URL
        """
        self.engine = create_engine(database_url)
        self.Session = sessionmaker(bind=self.engine)
    
    def write_gpu_usage(self, gpu_id, gpu_usage, gpu_memory_usage):
        """
        写入GPU使用情况到数据库
        :param gpu_id: GPU编号
        :param gpu_usage: GPU使用率（百分比）
        :param gpu_memory_usage: GPU内存用量（MB）
        :return: bool 是否写入成功
        """
        session = self.Session()
        try:
            gpu_record = GpuUsage(
                gpu_id=gpu_id,
                gpu_usage=gpu_usage,
                gpu_memory_usage=gpu_memory_usage
            )
            session.add(gpu_record)
            session.commit()
            return True
        except Exception as e:
            session.rollback()
            logging.error(f"写入GPU使用情况失败: {e}")
            return False
        finally:
            session.close()

    def get_latest_gpu_usage(self, limit=100):
        """
        获取最新的GPU使用情况记录
        :param limit: 获取记录条数，默认100条
        :return: list 包含GPU使用记录的字典列表
        """
        session = self.Session()
        try:
            records = session.query(GpuUsage)\
                           .order_by(GpuUsage.created_at.desc())\
                           .limit(limit)\
                           .all()
            
            result = []
            for record in records:
                result.append({
                    'id': record.gpu_id,  # 使用gpu_id作为id
                    'gpu_id': record.gpu_id,
                    'gpu_usage': record.gpu_usage,
                    'gpu_memory_usage': record.gpu_memory_usage,
                    'timestamp': record.created_at.isoformat() if record.created_at else None
                })
            
            return result
        except Exception as e:
            logging.error(f"获取GPU使用情况失败: {e}")
            return []
        finally:
            session.close()