Merge pull request #10458 from akx/graceful-stop

Graceful server stopping

1 files changed, 34 insertions, 16 deletions

diff --git a/webui.py b/webui.py
index 2cd551bd..cebfba96 100644
--- a/webui.py
+++ b/webui.py
@@ -8,7 +8,7 @@ import warnings
 import json

 from threading import Thread

 

-from fastapi import FastAPI

+from fastapi import FastAPI, Response

 from fastapi.middleware.cors import CORSMiddleware

 from fastapi.middleware.gzip import GZipMiddleware

 from packaging import version

@@ -241,7 +241,10 @@ def initialize():
         print(f'Interrupted with signal {sig} in {frame}')

         os._exit(0)

 

-    signal.signal(signal.SIGINT, sigint_handler)

+    if not os.environ.get("COVERAGE_RUN"):

+        # Don't install the immediate-quit handler when running under coverage,

+        # as then the coverage report won't be generated.

+        signal.signal(signal.SIGINT, sigint_handler)

 

 

 def setup_middleware(app):

@@ -262,19 +265,6 @@ def create_api(app):
     return api

 

 

-def wait_on_server(demo=None):

-    while 1:

-        time.sleep(0.5)

-        if shared.state.need_restart:

-            shared.state.need_restart = False

-            time.sleep(0.5)

-            demo.close()

-            time.sleep(0.5)

-

-            modules.script_callbacks.app_reload_callback()

-            break

-

-

 def api_only():

     initialize()

 

@@ -287,6 +277,12 @@ def api_only():
     print(f"Startup time: {startup_timer.summary()}.")

     api.launch(server_name="0.0.0.0" if cmd_opts.listen else "127.0.0.1", port=cmd_opts.port if cmd_opts.port else 7861)

 

+

+def stop_route(request):

+    shared.state.server_command = "stop"

+    return Response("Stopping.")

+

+

 def webui():

     launch_api = cmd_opts.api

     initialize()

@@ -335,6 +331,9 @@ def webui():
             inbrowser=cmd_opts.autolaunch,

             prevent_thread_lock=True

         )

+        if cmd_opts.add_stop_route:

+            app.add_route("/_stop", stop_route, methods=["POST"])

+

         # after initial launch, disable --autolaunch for subsequent restarts

         cmd_opts.autolaunch = False

 

@@ -366,8 +365,27 @@ def webui():
             redirector.get("/")

             gradio.mount_gradio_app(redirector, shared.demo, path=f"/{cmd_opts.subpath}")

 

-        wait_on_server(shared.demo)

+        try:

+            while True:

+                server_command = shared.state.wait_for_server_command(timeout=5)

+                if server_command:

+                    if server_command in ("stop", "restart"):

+                        break

+                    else:

+                        print(f"Unknown server command: {server_command}")

+        except KeyboardInterrupt:

+            print('Caught KeyboardInterrupt, stopping...')

+            server_command = "stop"

+

+        if server_command == "stop":

+            print("Stopping server...")

+            # If we catch a keyboard interrupt, we want to stop the server and exit.

+            shared.demo.close()

+            break

         print('Restarting UI...')

+        shared.demo.close()

+        time.sleep(0.5)

+        modules.script_callbacks.app_reload_callback()

 

         startup_timer.reset()