Colab Pro中PySpark训练时CPU/GPU/内存用量监控与存储问询

一、直接监控并写入TXT文件 #

适合需要本地存储监控数据的场景，通过Python脚本结合系统工具定时采集指标：

1. 安装依赖（若未安装）

!pip install psutil

2. 启动后台监控线程

import psutil
import time
from threading import Thread
import subprocess
import re

def monitor_resources(output_file="pyspark_resource_log.txt", interval=5):
    # 写入表头
    with open(output_file, "w") as f:
        f.write("timestamp,cpu_usage(%),memory_used(GB),memory_total(GB),gpu_usage(%),gpu_memory_used(GB),gpu_memory_total(GB)\n")
    
    while True:
        # 采集CPU与内存数据
        cpu_usage = psutil.cpu_percent(interval=1)
        mem = psutil.virtual_memory()
        mem_used = round(mem.used / (1024**3), 2)
        mem_total = round(mem.total / (1024**3), 2)
        
        # 采集GPU数据（Colab Pro专属）
        gpu_raw = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used,memory.total", "--format=csv,noheader,nounits"]
        ).decode("utf-8").strip()
        if gpu_raw:
            gpu_usage, gpu_mem_used, gpu_mem_total = [int(item.strip()) for item in gpu_raw.split(",")]
            gpu_mem_used = round(gpu_mem_used / 1024, 2)
            gpu_mem_total = round(gpu_mem_total / 1024, 2)
        else:
            gpu_usage, gpu_mem_used, gpu_mem_total = 0, 0, 0
        
        # 写入日志
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        with open(output_file, "a") as f:
            f.write(f"{timestamp},{cpu_usage},{mem_used},{mem_total},{gpu_usage},{gpu_mem_used},{gpu_mem_total}\n")
        
        time.sleep(interval - 1)

# 启动监控线程（后台运行，训练结束自动停止）
monitor_thread = Thread(target=monitor_resources, args=("pyspark_training_log.txt", 5))
monitor_thread.daemon = True
monitor_thread.start()

# 此处放置你的PySpark训练代码
# 示例：
# from pyspark.sql import SparkSession
# spark = SparkSession.builder.appName("ModelTraining").getOrCreate()
# ... 训练逻辑 ...

二、Weights & Biases (wandb) 完整使用教程 #

针对需要长期存储、可视化大量监控数据的场景，一步步实现集成：

1. 初始化wandb

!pip install wandb
import wandb
wandb.login()  # 按提示复制API密钥完成登录

2. 配置项目与监控线程

# 初始化wandb运行实例
wandb.init(
    project="pyspark-colab-training",
    name="run-001",
    config={"monitor_interval": 5}
)

def wandb_resource_monitor(interval=5):
    while True:
        # 采集资源数据
        cpu_usage = psutil.cpu_percent(interval=1)
        mem = psutil.virtual_memory()
        mem_used = round(mem.used / (1024**3), 2)
        
        gpu_raw = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used", "--format=csv,noheader,nounits"]
        ).decode("utf-8").strip()
        if gpu_raw:
            gpu_usage, gpu_mem_used = [int(item.strip()) for item in gpu_raw.split(",")]
            gpu_mem_used = round(gpu_mem_used / 1024, 2)
        else:
            gpu_usage, gpu_mem_used = 0, 0
        
        # 上传数据到wandb
        wandb.log({
            "cpu_usage(%)": cpu_usage,
            "memory_used(GB)": mem_used,
            "gpu_usage(%)": gpu_usage,
            "gpu_memory_used(GB)": gpu_mem_used
        })
        
        time.sleep(interval - 1)

# 启动监控线程
wandb_thread = Thread(target=wandb_resource_monitor, args=(5,))
wandb_thread.daemon = True
wandb_thread.start()

# 执行PySpark训练代码
# ... 训练逻辑 ...

# 训练结束后关闭wandb实例
wandb.finish()

3. 查看与导出数据

• 登录wandb控制台后，进入对应项目即可看到实时更新的资源监控图表
• 所有历史数据支持导出为CSV/TXT格式，满足后续分析需求
• 可扩展：在训练循环中加入wandb.log({"train_loss": loss_value})，实现训练指标与资源数据的统一监控

内容的提问来源于stack exchange，提问作者Iliasp