handle graceful exit when OOM,

and when the model couldn't be loaded

gunicorn.conf.py

@@ -2,4 +2,5 @@ workers = 1
 bind = "0.0.0.0:80"
 worker_class = "uvicorn.workers.UvicornWorker"
 timeout = 0
-keep_alive=60
\ No newline at end of file
+keep_alive=60
+pidfile="gunicorn.pid"
\ No newline at end of file

hydra_node/config.py

@@ -21,7 +21,7 @@ def init_config_model():
         config.amp = True
     elif config.amp == "0":
         config.amp = False
-        
+
     is_dev = ""
     environment = "production"
     if os.environ['DEV'] == "True":
@@ -80,8 +80,14 @@ def init_config_model():
     # Instantiate our actual model.
     load_time = time.time()
 
+    try:
+        model = StableDiffusionModel(config)
+    except Exception as e:
+        logger.error(f"Failed to load model: {str(e)}")
+        capture_exception(e)
+        #exit gunicorn
+        sys.exit(4)
 
-    model = StableDiffusionModel(config)
 
     config.model = model
 

main.py

@@ -20,8 +20,8 @@ import base64
 #Initialize model and config
 model, config = init_config_model()
 logger = config.logger
-#config.mainpid = open("app.pid", "r").read()
-
+config.mainpid = int(open("gunicorn.pid", "r").read())
+mainpid = config.mainpid
 hostname = socket.gethostname()
 sent_first_message = False
 
@@ -92,7 +92,7 @@ def generate(request: GenerationRequest):
         if "CUDA out of memory" in e_s or \
                 "an illegal memory access" in e_s or "CUDA" in e_s:
             logger.error("GPU error, committing seppuku.")
-            os.kill(os.getpid(), signal.SIGTERM)
+            os.kill(mainpid, signal.SIGTERM)
         return {"error": str(e)}
 
 if __name__ == "__main__":