Andrew/issue 383 fix hang

src/py/kaleido/__init__.py

@@ -12,25 +12,6 @@
 from ._page_generator import PageGenerator
 from .kaleido import Kaleido
 
-_global_server = _sync_server.GlobalKaleidoServer()
-
-
-def start_sync_server(*args, **kwargs):
-    """
-    Start a kaleido server which will process all sync generation requests.
-
-    Only one server can be started at a time.
-
-    This wrapper function takes the exact same arguments as kaleido.Kaleido().
-    """
-    _global_server.open(*args, **kwargs)
-
-
-def stop_sync_server():
-    """Stop the kaleido server. It can be restarted."""
-    _global_server.close()
-
-
 __all__ = [
     "Kaleido",
     "PageGenerator",
@@ -46,6 +27,40 @@ def stop_sync_server():
     "write_fig_sync",
 ]
 
+_global_server = _sync_server.GlobalKaleidoServer()
+
+
+def start_sync_server(*args, silence_warnings=False, **kwargs):
+    """
+    Start a kaleido server which will process all sync generation requests.
+
+    The kaleido server is a singleton, so it can't be opened twice. This
+    function will warn you if the server is already running.
+
+    This wrapper function takes the exact same arguments as kaleido.Kaleido(),
+    except one extra, `silence_warnings`.
+
+    Args:
+        *args: all arguments `Kaleido()` would take.
+        silence_warnings: (bool, default False): If True, don't emit warning if
+        starting an already started server.
+        **kwargs: all keyword arguments `Kaleido()` would take.
+
+    """
+    _global_server.open(*args, silence_warnings=silence_warnings, **kwargs)
+
+
+def stop_sync_server(*, silence_warnings=False):
+    """
+    Stop the kaleido server. It can be restarted. Warns if not started.
+
+    Args:
+        silence_warnings: (bool, default False): If True, don't emit warning if
+        stopping a server that's not running.
+
+    """
+    _global_server.close(silence_warnings=silence_warnings)
+
 
 async def calc_fig(
     fig,

src/py/kaleido/_sync_server.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import asyncio
+import atexit
 import warnings
+from functools import partial
 from queue import Queue
 from threading import Thread
 from typing import TYPE_CHECKING, NamedTuple
@@ -53,30 +55,38 @@ def __new__(cls):
     def is_running(self):
         return self._initialized
 
-    def open(self, *args: Any, **kwargs: Any) -> None:
+    def open(self, *args: Any, silence_warnings=False, **kwargs: Any) -> None:
         """Initialize the singleton with three values."""
         if self.is_running():
-            warnings.warn(
-                "Server already open.",
-                RuntimeWarning,
-                stacklevel=2,
-            )
+            if not silence_warnings:
+                warnings.warn(
+                    "Server already open.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
             return
         coroutine = self._server(*args, **kwargs)
-        self._thread: Thread = Thread(target=asyncio.run, args=(coroutine,))
+        self._thread: Thread = Thread(
+            target=asyncio.run,
+            args=(coroutine,),
+            daemon=True,
+        )
         self._task_queue: Queue[Task | None] = Queue()
         self._return_queue: Queue[Any] = Queue()
         self._thread.start()
         self._initialized = True
+        close = partial(self.close, silence_warnings=True)
+        atexit.register(close)
 
-    def close(self):
+    def close(self, *, silence_warnings=False):
         """Reset the singleton back to an uninitialized state."""
         if not self.is_running():
-            warnings.warn(
-                "Server already closed.",
-                RuntimeWarning,
-                stacklevel=2,
-            )
+            if not silence_warnings:
+                warnings.warn(
+                    "Server already closed.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
             return
         self._task_queue.put(None)
         self._thread.join()