From 5a1b5d3de408f80ca8a1c714456f15e25309c6bd Mon Sep 17 00:00:00 2001
From: Guokan SHANG <guokan.shang@hotmail.com>
Date: Wed, 12 Apr 2023 14:15:20 +0200
Subject: [PATCH] gpu

---
 .envdefault            |  1 +
 docker-compose.yml     |  2 +-
 docker-entrypoint.sh   |  3 ++-
 http_server/ingress.py | 14 +++++++++++---
 http_server/serving.py | 24 +++++++++++++++++++++++-
 5 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/.envdefault b/.envdefault
index bd2e9a5..e3d950c 100644
--- a/.envdefault
+++ b/.envdefault
@@ -10,6 +10,7 @@ LM_MAP = "{
 # SERVING PARAMETERS
 SERVICE_MODE=http
 CONCURRENCY=1
+USE_GPU=True
 
 # MICRO-SERVICE PARAMETERS
 SERVICE_NAME=kpe
diff --git a/docker-compose.yml b/docker-compose.yml
index 5f5a823..97e09f6 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -9,4 +9,4 @@ services:
       - .env
     volumes:
       - $ASSETS_PATH_ON_HOST:$ASSETS_PATH_IN_CONTAINER:ro
-    #runtime: nvidia
\ No newline at end of file
+    runtime: nvidia
\ No newline at end of file
diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh
index e4bb9aa..a697890 100755
--- a/docker-entrypoint.sh
+++ b/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
diff --git a/http_server/ingress.py b/http_server/ingress.py
index 5181508..4f35b03 100644
--- a/http_server/ingress.py
+++ b/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:
diff --git a/http_server/serving.py b/http_server/serving.py
index 076f34d..68d32e5 100644
--- a/http_server/serving.py
+++ b/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
\ No newline at end of file
+        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()
\ No newline at end of file