gunicorn下的深度学习api 如何合理分配gpu

栏目: Python · 发布时间: 6年前

内容简介:老大提了一个需求: gunicron 起多个进程的时候,如何保证pytorch的模型均匀分配到不同的gpu上,按道理,如果能拿到类似每个进程的序号,那分配起来应该都是简单的,那核心问题提炼出来了,如何拿到进程的序号顺手直接去找一个相关的问题和分析,https://github.com/benoitc/gunicorn/issues/1278 ,发现很多人都有同样的需求,不过貌似提的pr都没有进一步的解决,所以只能进一步来看官方的文档有什么可用的。

背景

老大提了一个需求: gunicron 起多个进程的时候,如何保证pytorch的模型均匀分配到不同的gpu上,按道理,如果能拿到类似每个进程的序号,那分配起来应该都是简单的,那核心问题提炼出来了,如何拿到进程的序号

分析

顺手直接去找一个相关的问题和分析,https://github.com/benoitc/gunicorn/issues/1278 ,发现很多人都有同样的需求,不过貌似提的pr都没有进一步的解决,所以只能进一步来看官方的文档有什么可用的。

gunicorn下的深度学习api 如何合理分配gpu

通过进一步发现 http://docs.gunicorn.org/en/latest/settings.html 的文档,这些在起进程的时候就可以预先定义好进程的id

实践

我们写好 gunicorn_conf.py

# RTFM -> http://docs.gunicorn.org/en/latest/settings.html#settings
import os
 
from service.config import WORKERS
 
bind = '0.0.0.0:2048'
workers = WORKERS
timeout = 300
max_requests = 2000
max_requests_jitter = 500
 
 
def on_starting(server):
    """
    Attach a set of IDs that can be temporarily re-used.
    Used on reloads when each worker exists twice.
    """
    server._worker_id_overload = set()
 
 
def nworkers_changed(server, new_value, old_value):
    """
    Gets called on startup too.
    Set the current number of workers.  Required if we raise the worker count
    temporarily using TTIN because server.cfg.workers won't be updated and if
    one of those workers dies, we wouldn't know the ids go that far.
    """
    server._worker_id_current_workers = new_value
 
 
def _next_worker_id(server):
    """
    If there are IDs open for re-use, take one.  Else look for a free one.
    """
    if server._worker_id_overload:
        return server._worker_id_overload.pop()
 
    in_use = set(w._worker_id for w in server.WORKERS.values() if w.alive)
    free = set(range(1, server._worker_id_current_workers + 1)) - in_use
 
    return free.pop()
 
 
def on_reload(server):
    """
    Add a full set of ids into overload so it can be re-used once.
    """
    server._worker_id_overload = set(range(1, server.cfg.workers + 1))
 
 
def pre_fork(server, worker):
    """
    Attach the next free worker_id before forking off.
    """
    worker._worker_id = _next_worker_id(server)
 
 
def post_fork(server, worker):
    """
    Put the worker_id into an env variable for further use within the app.
    """
    os.environ["APP_WORKER_ID"] = str(worker._worker_id)
 
 

这样我们通过环境变量就可以清楚的知道我们的当前子进程的序号

# -*- coding: utf-8 -*-
import os
 
import torch
 
 
def set_process_gpu():
    worker_id = int(os.environ.get('APP_WORKER_ID', 1))
    devices = os.environ.get('CUDA_VISIBLE_DEVICES', '')
 
    if not devices:
        print('current environment did not get CUDA_VISIBLE_DEVICES env ,so use the default')
 
    rand_max = 9527
    gpu_index = (worker_id + rand_max) % torch.cuda.device_count()
 
    print('current worker id  {} set the gpu id :{}'.format(worker_id, gpu_index))
    torch.cuda.set_device(int(gpu_index))
 

通过这个方法就可以轻松的设置自己进程所在的gpu ,这样就可以根据gpu的数量,均匀的分配进程

gunicorn -c gunicorn_conf.py wsgi:app
 

wsgi.py 这个就是app的实体了,正常启用就可以了。


以上就是本文的全部内容,希望对大家的学习有所帮助,也希望大家多多支持 码农网

查看所有标签

猜你喜欢:

本站部分资源来源于网络,本站转载出于传递更多信息之目的,版权归原作者或者来源机构所有,如转载稿涉及版权问题,请联系我们

怎样解题

怎样解题

[美] G. 波利亚 / 涂泓、冯承天 / 上海科技教育出版社 / 2002-6 / 16.00元

《怎样解题:数学教学法的新面貌》是数学家波利亚论述中学数学教学法的普及名著,对数学教育产生了深刻的影响。波利亚认为中学数学教育的根本宗旨是教会年轻人思考,他把“解题”作为培养学生数学才能和教会他们思考的一种手段和途径。这本书是他专门研究解题的思维过程后的结晶。全书的核心是他分解解题的思维过程得到的一张“怎样解题”表。作者在书中引导学生按照“表”中的问题和建议思考问题,探索解题途径,进而逐步掌握解题......一起来看看 《怎样解题》 这本书的介绍吧!

CSS 压缩/解压工具
CSS 压缩/解压工具

在线压缩/解压 CSS 代码

Base64 编码/解码
Base64 编码/解码

Base64 编码/解码

XML 在线格式化
XML 在线格式化

在线 XML 格式化压缩工具