handle graceful exit when OOM,

and when the model couldn't be loaded

gunicorn.conf.py

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

hydra_node/config.py

@@ -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)
+
 
     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__":