Update jupyter/jupyter_server.py

jupyter/jupyter_server.py

@@ -1,4 +1,5 @@
 from flask import Flask, request, jsonify
+import os
 import uuid
 import time
 import docker
@@ -9,6 +10,7 @@ import argparse
 import logging
 from pydantic import BaseModel, Field, ValidationError
 
+current_dir = os.path.dirname(os.path.abspath(__file__))
 
 app = Flask(__name__)
 app.logger.setLevel(logging.INFO)
@@ -44,7 +46,7 @@ def create_kernel_containers(n_instances, n_cpus=2, mem="2g", execution_timeout=
 
     docker_client = docker.from_env()
     app.logger.info("Buidling docker image...")
-    image, logs = docker_client.images.build(path='./', tag='jupyter-kernel:latest')
+    image, logs = docker_client.images.build(path=current_dir, tag='jupyter-kernel:latest')
     app.logger.info("Building docker image complete.")
 
     containers = []
@@ -139,18 +141,33 @@ def execute_code():
         }), 500
 
 
-atexit.register(shutdown_cleanup)
+def init_app(app, args=None):
+    if args is None:
+        # When run through Gunicorn, use environment variables
+        args = argparse.Namespace(
+            n_instances=int(os.getenv('N_INSTANCES', 1)),
+            n_cpus=int(os.getenv('N_CPUS', 1)),
+            mem=os.getenv('MEM', '1g'),
+            execution_timeout=int(os.getenv('EXECUTION_TIMEOUT', 60))
+        )
 
-if __name__ == '__main__':
-    args = parse_args()
     app.containers = create_kernel_containers(
         args.n_instances,
         n_cpus=args.n_cpus, 
         mem=args.mem, 
         execution_timeout=args.execution_timeout
-        )
+    )
+    return app, args
+
+atexit.register(shutdown_cleanup)
+
+if __name__ == '__main__':
+    args = parse_args()
+    app, args = init_app(app, args=args)
     # don't use debug=True --> it will run main twice and thus start double the containers
     app.run(debug=False, host='0.0.0.0', port=args.port) 
+else:
+    app, args = init_app(app)
 
 
 # TODO: