import psutil
import time
from datetime import datetime
from pynvml import *

def get_memory_usage():
    mem = psutil.virtual_memory()
    return mem.percent, mem.used / (1024 ** 3), mem.total / (1024 ** 3)

def get_cpu_usage():
    return psutil.cpu_percent(interval=1)  # 获取CPU利用率（百分比）

def get_gpu_usage():
    try:
        nvmlInit()
        device_count = nvmlDeviceGetCount()
        gpu_stats = []
        for i in range(device_count):
            handle = nvmlDeviceGetHandleByIndex(i)
            name = nvmlDeviceGetName(handle).decode("utf-8")
            mem_info = nvmlDeviceGetMemoryInfo(handle)
            util = nvmlDeviceGetUtilizationRates(handle)
            gpu_stats.append({
                "name": name,
                "gpu_util": util.gpu,
                "mem_util": util.memory,
                "mem_used": mem_info.used / (1024 ** 3),
                "mem_total": mem_info.total / (1024 ** 3)
            })
        nvmlShutdown()
        return gpu_stats
    except Exception:
        return [{"name": "N/A", "gpu_util": 0, "mem_util": 0, "mem_used": 0, "mem_total": 0}]

def main():
    log_file = "monitor_log.txt"
    with open(log_file, "a", encoding="utf-8") as f:
        while True:
            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            
            # CPU
            cpu_percent = get_cpu_usage()
            
            # 内存
            mem_percent, mem_used, mem_total = get_memory_usage()
            
            # GPU
            gpu_stats = get_gpu_usage()
            
            # 拼接日志内容
            log_line = (f"[{now}] CPU: {cpu_percent:.1f}% | "
                        f"Memory: {mem_percent:.1f}% ({mem_used:.2f} GB / {mem_total:.2f} GB)")
            
            for gpu in gpu_stats:
                log_line += (f" | GPU: {gpu['name']} {gpu['gpu_util']}% "
                             f"(Mem {gpu['mem_used']:.2f} GB / {gpu['mem_total']:.2f} GB)")
            
            # 输出
            print(log_line)
            f.write(log_line + "\n")
            f.flush()
            
            time.sleep(9)  # 这里+CPU统计的1秒，总体大约10秒一个周期

if __name__ == "__main__":
    main()