[工程] gunicorn下的深度学习api 如何合理分配gpu

背景

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

分析

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

通过进一步发现 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)

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

# 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))

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

# -*- 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

1

2

gunicorn -c gunicorn_conf.py wsgi:app

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

API GPU加速云服务器 深度学习

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。