创建两个独立的队列

为每个GPU创建一个单独的队列配置文件:

1
2
3
mkdir -p ~/.config/task-spooler
touch ~/.config/task-spooler/gpu0
touch ~/.config/task-spooler/gpu1

设置队列的GPU分配

编辑队列配置文件,为每个队列指定使用的GPU:

1
2
3
4
5
# 配置GPU0队列
echo "export CUDA_VISIBLE_DEVICES=0" > ~/.config/task-spooler/gpu0

# 配置GPU1队列
echo "export CUDA_VISIBLE_DEVICES=1" > ~/.config/task-spooler/gpu1

启动两个独立的Task Spooler服务器

1
2
3
4
5
# 启动GPU0队列的服务器
TS_SOCKET=~/.ts_gpu0 tsp -S

# 启动GPU1队列的服务器
TS_SOCKET=~/.ts_gpu1 tsp -S

向不同的队列提交任务

使用TS_SOCKET环境变量指定队列,向不同GPU提交任务:

1
2
3
4
5
# 向GPU0队列提交任务
TS_SOCKET=~/.ts_gpu0 tsp python my_script.py --param1 value1

# 向GPU1队列提交任务
TS_SOCKET=~/.ts_gpu1 tsp python my_script.py --param2 value2

查看队列状态

查看各个队列的状态

1
2
3
4
5
# 查看GPU0队列状态
TS_SOCKET=~/.ts_gpu0 tsp

# 查看GPU1队列状态
TS_SOCKET=~/.ts_gpu1 tsp

自动化配置

为了方便使用,可以将一下别名添加到~/.bashrc中:

1
2
3
4
5
6
7
# Task Spooler GPU队列别名
alias tsp0='TS_SOCKET=~/.ts_gpu0 tsp'
alias tsp1='TS_SOCKET=~/.ts_gpu1 tsp'

# 启动GPU队列服务器的别名
alias tsp0-start='TS_SOCKET=~/.ts_gpu0 tsp -S'
alias tsp1-start='TS_SOCKET=~/.ts_gpu1 tsp -S'