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

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

gunicorn -c gunicorn_conf.py wsgi:app