gpu

Author: guokan-shang

.envdefault

@@ -10,6 +10,7 @@ LM_MAP = "{
 # SERVING PARAMETERS
 SERVICE_MODE=http
 CONCURRENCY=1
+USE_GPU=True
 
 # MICRO-SERVICE PARAMETERS
 SERVICE_NAME=kpe

docker-compose.yml

@@ -9,4 +9,4 @@ services:
  - .env
  volumes:
  - $ASSETS_PATH_ON_HOST:$ASSETS_PATH_IN_CONTAINER:ro
- #runtime: nvidia
+ runtime: nvidia

docker-entrypoint.sh

@@ -22,7 +22,8 @@ else
  fi
  /usr/src/app/wait-for-it.sh $(echo $SERVICES_BROKER | cut -d'/' -f 3) --timeout=20 --strict -- echo " $SERVICES_BROKER (Service Broker) is up"
  echo "RUNNING CELERY WORKER"
- celery --app=celery_app.celeryapp worker -Ofair -n nlp_${SERVICE_NAME}_worker@%h --queues=${SERVICE_NAME} -c ${CONCURRENCY}
+ POOL=$([ $USE_GPU == "True" ] && echo "gevent" || echo "prefork")
+ celery --app=celery_app.celeryapp worker -Ofair -n nlp_${SERVICE_NAME}_worker@%h --queues=${SERVICE_NAME} -c ${CONCURRENCY} --pool=$POOL
  else
  echo "ERROR: Wrong serving command: $1"
  exit -1

http_server/ingress.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import os
 import json
 import logging
 from time import time
@@ -8,7 +9,7 @@
 import components
 
 from flask import Flask, request, abort, Response, json
-from serving import GunicornServing
+from serving import GeventServing, GunicornServing
 from confparser import createParser
 from swagger import setupSwaggerUI
 
@@ -86,8 +87,15 @@ def server_error(error):
  logger.debug("Swagger UI set.")
  except Exception as e:
  logger.warning("Could not setup swagger: {}".format(str(e)))
- 
- serving = GunicornServing(app, {'bind': '{}:{}'.format("0.0.0.0", args.service_port),
+
+ if os.environ.get("USE_GPU", "True") == "True":
+ serving_type = GeventServing
+ logger.debug("Serving with gevent")
+ else:
+ serving_type = GunicornServing
+ logger.debug("Serving with gunicorn")
+
+ serving = serving_type(app, {'bind': '{}:{}'.format("0.0.0.0", args.service_port),
  'workers': args.workers,})
  logger.info(args)
  try:

http_server/serving.py

@@ -1,4 +1,7 @@
 import gunicorn.app.base
+import gevent.pywsgi
+import gevent.monkey
+gevent.monkey.patch_all()
 
 class GunicornServing(gunicorn.app.base.BaseApplication):
 
@@ -14,4 +17,23 @@ def load_config(self):
  self.cfg.set(key.lower(), value)
 
  def load(self):
- return self.application
+ return self.application
+
+class GeventServing():
+
+ def __init__(self, app, options=None):
+ self.options = options or {}
+ self.application = app
+
+ def run(self):
+ bind = self.options.get('bind', "0.0.0.0:8080")
+ workers = self.options.get('workers', 1)
+ listener = bind.split(':')
+ try:
+ assert len(listener) == 2
+ listener = (listener[0], int(listener[1]))
+ except:
+ print(f"Invalid bind address {bind}")
+
+ server = gevent.pywsgi.WSGIServer(listener, self.application, spawn = workers)
+ server.serve_forever()