gpu

.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()