feat(api): switch API server to Waitress (#159)

ssube · Apr 11, 2023 · 76860f6 · 76860f6
1 parent e3540ac
commit 76860f6
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 6 deletions.
diff --git a/api/launch-extras.bat b/api/launch-extras.bat
@@ -9,4 +9,10 @@ python -m onnx_web.convert ^
 --token=%HF_TOKEN% %ONNX_WEB_EXTRA_ARGS%
 
 echo "Launching API server..."
-flask --app="onnx_web.main:run" run --host=0.0.0.0
+waitress-serve ^
+--host=0.0.0.0 ^
+--port=5000 ^
+--call ^
+onnx_web.main:run
+
+pause
diff --git a/api/launch-extras.sh b/api/launch-extras.sh
@@ -26,4 +26,8 @@ python3 -m onnx_web.convert \
   ${ONNX_WEB_EXTRA_ARGS:-}
 
 echo "Launching API server..."
-flask --app='onnx_web.main:run' run --host=0.0.0.0
+waitress-serve \
+  --host=0.0.0.0 \
+  --port=5000 \
+  --call \
+  onnx_web.main:run
diff --git a/api/launch.bat b/api/launch.bat
@@ -7,4 +7,10 @@ python -m onnx_web.convert ^
 --token=%HF_TOKEN% %ONNX_WEB_EXTRA_ARGS%
 
 echo "Launching API server..."
-flask --app="onnx_web.main:run" run --host=0.0.0.0
+waitress-serve ^
+--host=0.0.0.0 ^
+--port=5000 ^
+--call ^
+onnx_web.main:run
+
+pause
diff --git a/api/launch.sh b/api/launch.sh
@@ -25,4 +25,8 @@ python3 -m onnx_web.convert \
   ${ONNX_WEB_EXTRA_ARGS:-}
 
 echo "Launching API server..."
-flask --app='onnx_web.main:run' run --host=0.0.0.0
+waitress-serve \
+  --host=0.0.0.0 \
+  --port=5000 \
+  --call \
+  onnx_web.main:run
diff --git a/api/onnx_web/main.py b/api/onnx_web/main.py
@@ -78,7 +78,9 @@ def quit(p: DevicePoolExecutor):
 
 if __name__ == "__main__":
     app, pool = main()
+    logger.info("starting image workers")
     pool.start()
+    logger.info("starting API server")
     app.run("0.0.0.0", 5000, debug=is_debug())
-    logger.info("shutting down app")
+    logger.info("shutting down workers")
     pool.join()
diff --git a/api/requirements/base.txt b/api/requirements/base.txt
@@ -25,4 +25,5 @@ flask
 flask-cors
 jsonschema
 pyyaml
-setproctitle
+setproctitle
+waitress