如何优雅地抢占显卡

各位是否一直头疼于服务器抢不上卡？不用担心，在我的邀请下，deepseek给大家写了一个自动检测显卡空闲情况并及时抢占的代码。使用只需要简单改部分参数即可~

需要实现的功能如下：

• 检测显卡状态
• 在终端执行需要跑的命令
• 每隔一段时间进行检测

这些功能下面都实现了，本人亲测能跑，需要修改的部分有

• 显卡数，如果你只需要1张卡，在main函数的if_available>=8那里把8改成1即可（其他同理）
• 时间啥的自己改
• 激活conda环境名，自己改自己的就好，以及conda.sh位置自己找一下改了就行

如下：

def run_training():
    """执行训练命令"""
    try:
        cmd = [
            "/bin/bash", "-c",
            "source /xxx/miniconda3/etc/profile.d/conda.sh && "
            "conda activate xxx && "
            "FORCE_TORCHRUN=1 llamafactory-cli train "
            "/xxx/LLaMA-Factory/examples/train_full/qwen2_5vl_full_sft.yaml"
        ]
        subprocess.run(cmd, check=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"训练任务启动失败: {e}")
        return False

这里面，source /xxx/miniconda3/etc/profile.d/conda.sh 就是在准备conda环境（具体叫啥忘了，理解为前期准备），执行完这个就可以用conda activate xxx激活conda环境了。

下面是最终代码，nvidia-smi后面那一串不同机器可能报不同的错，自己找deepseek捣鼓捣鼓就行。

import subprocess
import time
import sys

def get_available_gpus():
    """获取前8张显卡的空闲状态"""
    try:
        # 使用通用字段查询GPU使用情况
        cmd = [
            "nvidia-smi",
            "--query-gpu=index,utilization.gpu,memory.used,memory.total",
            "--format=csv,noheader,nounits"
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        
        available_gpus = 0
        for line in result.stdout.strip().split('\n'):
            if not line:
                continue
            
            # 解析GPU信息
            gpu_id, gpu_util, mem_used, mem_total = map(int, line.split(', '))
            
            # 仅检查前8张卡
            if gpu_id >= 8:
                continue
                
            # 空闲判断标准（可调整阈值）
            if gpu_util < 5 and (mem_used / mem_total) < 0.1:
                available_gpus += 1

        return available_gpus
    
    except Exception as e:
        print(f"GPU状态检测失败: {e}")
        return 0

def run_training():
    """执行训练命令"""
    try:
        cmd = [
            "/bin/bash", "-c",
            "source /xxx/miniconda3/etc/profile.d/conda.sh && "
            "conda activate xxx && "
            "FORCE_TORCHRUN=1 llamafactory-cli train "
            "/xxx/LLaMA-Factory/examples/train_full/qwen2_5vl_full_sft.yaml"
        ]
        subprocess.run(cmd, check=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"训练任务启动失败: {e}")
        return False

def main():
    print("🛠️ GPU监控系统启动...")
    check_interval = 300  # 默认20分钟检测一次
    
    while True:
        try:
            available = get_available_gpus()
            print(f"▏ 可用显卡：{available}/8 | 下次检测：{check_interval//60}分钟后")
            
            if available >= 8:
                print("🚦 满足条件，启动训练进程...")
                if run_training():
                    print("✅ 训练任务已部署")
                    sys.exit(0)
                else:
                    check_interval = 1800  # 失败后30分钟重试
            else:
                # 动态调整检测间隔（1-15分钟）
                check_interval = min(900, max(60, (8 - available) * 120))
            
            time.sleep(check_interval)
            
        except KeyboardInterrupt:
            print("\n🔌 用户终止监控进程")
            sys.exit(1)

if __name__ == "__main__":
    main()

OK，就到这里，祝大家不用再蹲GPU，都能优雅地度过每一个跑实验的日子…