前言

你是否遇到过这样的场景？

• 实验室有一台 8 卡 GPU 服务器，几个同学同时跑实验，经常互相 "抢"GPU (是的，我们组是穷组)
• 想跑一个超参数搜索，需要提交 10 个不同的实验，还要记得哪个实验用了哪些参数
• 训练完模型后需要做评估，但总是要手动盯着训练什么时候结束

Gflow 就是这样一个解决方案。

Gflow 是一个用 Rust 编写的轻量单节点作业调度器，专为管理和调度机器学习 / 深度学习任务而设计，特别是在具有多 GPU 资源的机器上。

他会帮你：

• 自动排队：提交任务后自动排队，有空闲 GPU 就运行
• 时间限制：防止失控任务无限占用资源
• 任务依赖：自动管理多阶段流水线（数据处理→训练→评估）
• 实时监控：随时查看任务状态和资源使用情况
• 日志管理：自动保存每个任务的输出日志

安装

gflow 支持 cargo, conda 安装

conda install -c conda-forge gflow
# or
cargo install gflow

快速上手

# 启动后端
gflowd up

# 提交需要 1 个 GPU 的训练任务
gbatch --gpus 1 python train.py
# 依赖上一个任务结束
gbatch --depends-on @~1 python evaluate.py 

查看依赖关系树：

gqueue -t

# 输出：
JOBID    NAME         ST    TIME         TIMELIMIT
1        preprocess   CD    00:02:15     00:20:00
└─ 2     train        R     00:15:30     04:00:00
   └─ 3  evaluate     PD    00:00:00     00:10:00

搜索参数

gbatch -g 1 --param scale=2:0:-0.1 python train.py --model_id 0 --scale {scale}

自动探测 conda 环境

在默认情况下，gbatch 将会探测当前 shell 激活的环境，并作为脚本执行前激活的 conda 环境

conda activate lala
gbatch -g 1 --param scale=2:0:-0.1 python train.py --model_id 0 --scale {scale}
# or 手动设置
gbatch -c tran -g 1 --param scale=2:0:-0.1 python train.py --model_id 0 --scale {scale}

https://github.com/AndPuQing/gflow