solara-ui 1.31.0__py2.py3-none-any.whl

solara/server/kernel.py

@@ -0,0 +1,338 @@
+import json
+import logging
+import pdb
+import queue
+import struct
+import warnings
+from binascii import b2a_base64
+from datetime import datetime
+from typing import Set, Union
+
+import ipykernel
+import ipykernel.kernelbase
+import jupyter_client.session as session
+from dateutil.tz import tzlocal  # type: ignore
+from ipykernel.comm import CommManager
+from zmq.eventloop.zmqstream import ZMQStream
+
+import solara
+from solara.server.shell import SolaraInteractiveShell
+
+from . import settings, websocket
+
+logger = logging.getLogger("solara.server.kernel")
+ipykernel_major = int(ipykernel.__version__.split(".")[0])
+
+jsonmodule = json
+
+
+# from jupyter_client/jsonutil.py
+def _ensure_tzinfo(dt: datetime) -> datetime:
+    """Ensure a datetime object has tzinfo
+    If no tzinfo is present, add tzlocal
+    """
+    if not dt.tzinfo:
+        # No more naïve datetime objects!
+        warnings.warn(
+            "Interpreting naive datetime as local %s. Please add timezone info to timestamps." % dt,
+            DeprecationWarning,
+            stacklevel=4,
+        )
+        dt = dt.replace(tzinfo=tzlocal())
+    return dt
+
+
+def json_default(obj):
+    """default function for packing objects in JSON."""
+    if isinstance(obj, datetime):
+        obj = _ensure_tzinfo(obj)
+        return obj.isoformat().replace("+00:00", "Z")
+    elif isinstance(obj, bytes):
+        return b2a_base64(obj).decode("ascii")
+    if type(obj).__module__ == "numpy":
+        import numpy as np
+
+        if isinstance(obj, np.number):
+            return repr(obj.item())
+        else:
+            raise TypeError("%r is not JSON serializable" % obj)
+    else:
+        raise TypeError("%r is not JSON serializable" % obj)
+
+
+def json_dumps(data):
+    try:
+        return jsonmodule.dumps(data)
+    except TypeError:
+        logger.warning("Invalid JSON, will try a with a more forgiving json encoder")
+    return jsonmodule.dumps(
+        data,
+        default=json_default,
+        ensure_ascii=False,
+        allow_nan=False,
+    )
+
+
+ipykernel_version = tuple(map(int, ipykernel.__version__.split(".")))
+if ipykernel_version >= (6, 18, 0):
+    import comm.base_comm
+
+    class Comm(comm.base_comm.BaseComm):
+        kernel: Union[ipykernel.kernelbase.Kernel, None]
+
+        def __init__(self, **kwargs) -> None:
+            if ipykernel.kernelbase.Kernel.initialized():
+                self.kernel = ipykernel.kernelbase.Kernel.instance()
+            else:
+                self.kernel = None
+            super().__init__(**kwargs)
+
+        def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
+            if self.kernel is None or self.kernel.session is None:
+                return
+            data = {} if data is None else data
+            metadata = {} if metadata is None else metadata
+            content = dict(data=data, comm_id=self.comm_id, **keys)
+            self.kernel.session.send(
+                self.kernel.iopub_socket,
+                msg_type,
+                content,
+                metadata=metadata,
+                parent=self.kernel.get_parent("shell"),
+                ident=self.topic,
+                buffers=buffers,
+            )
+
+    comm.create_comm = Comm
+
+    def get_comm_manager():
+        from .kernel_context import get_current_context, has_current_context
+
+        if has_current_context():
+            return get_current_context().kernel.comm_manager
+        else:
+            return global_comm_manager
+
+    global_comm_manager = comm.get_comm_manager()
+    comm.get_comm_manager = get_comm_manager
+
+# from notebook.base.zmqhandlers import serialize_binary_message
+# this saves us a dependency on notebook/jupyter_server when e.g.
+# running on pyodide
+
+
+def _fix_msg(msg):
+    # makes sure the msg can be json serializable
+    # instead of using a callable like in jupyter_client (i.e. json_default)
+    # we replace the keys we know are problematic
+    # this allows us to use a faster json serializer in the future
+    if "header" in msg and "date" in msg["header"]:
+        # this is what jupyter_client.jsonutil.json_default does
+        msg["header"]["date"] = msg["header"]["date"].isoformat().replace("+00:00", "Z")
+    if "parent_header" in msg and "date" in msg["parent_header"]:
+        # date is already a string if it's copied from the header that is not turned into a datetime
+        # maybe we should do that in server.py
+        date = msg["parent_header"]["date"]
+        if isinstance(date, datetime):
+            msg["parent_header"]["date"] = date.isoformat().replace("+00:00", "Z")
+
+
+def serialize_binary_message(msg):
+    """serialize a message as a binary blob
+
+    Header:
+
+    4 bytes: number of msg parts (nbufs) as 32b int
+    4 * nbufs bytes: offset for each buffer as integer as 32b int
+
+    Offsets are from the start of the buffer, including the header.
+
+    Returns
+    -------
+
+    The message serialized to bytes.
+
+    """
+    # don't modify msg or buffer list in-place
+    msg = msg.copy()
+    buffers = list(msg.pop("buffers"))
+    bmsg = json_dumps(msg).encode("utf8")
+    buffers.insert(0, bmsg)
+    nbufs = len(buffers)
+    offsets = [4 * (nbufs + 1)]
+    for buf in buffers[:-1]:
+        offsets.append(offsets[-1] + len(buf))
+    offsets_buf = struct.pack("!" + "I" * (nbufs + 1), nbufs, *offsets)
+    buffers.insert(0, offsets_buf)
+    return b"".join(buffers)
+
+
+def deserialize_binary_message(bmsg):
+    """deserialize a message from a binary blog
+
+    Header:
+
+    4 bytes: number of msg parts (nbufs) as 32b int
+    4 * nbufs bytes: offset for each buffer as integer as 32b int
+
+    Offsets are from the start of the buffer, including the header.
+
+    Returns
+    -------
+    message dictionary
+    """
+    nbufs = struct.unpack("!i", bmsg[:4])[0]
+    offsets = list(struct.unpack("!" + "I" * nbufs, bmsg[4 : 4 * (nbufs + 1)]))
+    offsets.append(None)
+    bufs = []
+    for start, stop in zip(offsets[:-1], offsets[1:]):
+        bufs.append(bmsg[start:stop])
+    msg = json.loads(bufs[0].decode("utf8"))
+    msg["buffers"] = bufs[1:]
+    return msg
+
+
+SESSION_KEY = b"solara"
+
+
+class WebsocketStream:
+    def __init__(self, session, channel: str):
+        self.session = session
+        self.channel = channel
+
+
+class WebsocketStreamWrapper(ZMQStream):
+    def __init__(self, websocket, channel):
+        self.websocket = websocket
+        self.channel = channel
+
+    def flush(self, *ignore):
+        pass
+
+
+def send_websockets(websockets: Set[websocket.WebsocketWrapper], binary_msg):
+    for ws in list(websockets):
+        try:
+            ws.send(binary_msg)
+        except:  # noqa
+            # in case of any issue, we simply remove it from the list
+            try:
+                # websocket can be modified by another thread
+                websockets.remove(ws)
+            except KeyError:
+                pass  # already removed
+
+
+class SessionWebsocket(session.Session):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.websockets: Set[websocket.WebsocketWrapper] = set()  # map from .. msg id to websocket?
+
+    def close(self):
+        for ws in list(self.websockets):
+            try:
+                ws.close()
+            except:  # noqa
+                pass
+
+    def send(
+        self,
+        stream,
+        msg_or_type,
+        content=None,
+        parent=None,
+        ident=None,
+        buffers=None,
+        track=False,
+        header=None,
+        metadata=None,
+    ):
+        try:
+            if isinstance(msg_or_type, dict):
+                msg = msg_or_type
+            else:
+                msg = self.msg(
+                    msg_or_type,
+                    content=content,
+                    parent=parent,
+                    header=header,
+                    metadata=metadata,
+                )
+            _fix_msg(msg)
+            msg["channel"] = stream.channel
+            # not using pdb guard for performance reasons
+            try:
+                if buffers:
+                    msg["buffers"] = [memoryview(k).cast("b") for k in buffers]
+                    wire_message = serialize_binary_message(msg)
+                else:
+                    wire_message = json_dumps(msg)
+            except Exception:
+                logger.exception("Could not serialize message: %r", msg)
+                if settings.main.use_pdb:
+                    pdb.post_mortem()
+                raise
+            send_websockets(self.websockets, wire_message)
+        except Exception as e:
+            logger.exception("Error sending message: %s", e)
+
+
+class Kernel(ipykernel.kernelbase.Kernel):
+    session: SessionWebsocket  # type: ignore
+    # Ideally we have `session = Instance(Session, allow_none=True)`, but MyPy does not like it
+
+    implementation = "solara"
+    implementation_version = solara.__version__
+    banner = "solara"
+
+    def __init__(self):
+        super().__init__()
+        self.session = SessionWebsocket(parent=self, key=SESSION_KEY)
+        self.msg_queue = queue.Queue()  # type: ignore
+        self.stream = self.iopub_socket = WebsocketStream(self.session, "iopub")
+        # on github action the next line gives a mypy error:
+        # solara/server/kernel.py:111: error: "SessionWebsocket" has no attribute "stream"
+        # not sure why we cannot reproduce that locally
+        self.session.stream = self.iopub_socket  # type: ignore
+        if ipykernel_version >= (6, 18, 0):
+            # from this version on, ipykernel uses the comm package https://github.com/ipython/ipykernel/pull/973
+            self.comm_manager = CommManager(parent=self, kernel=self)
+            import ipywidgets.widgets.widget
+
+            if hasattr(ipywidgets.widgets.widget, "Comm"):
+                ipywidgets.widgets.widget.Comm = Comm
+                ipywidgets.widgets.widget.Widget.comm.klass = Comm
+        else:
+            self.comm_manager = CommManager(parent=self, kernel=self)
+        self.log = logging.getLogger("fake")
+
+        comm_msg_types = ["comm_open", "comm_msg", "comm_close"]
+        for msg_type in comm_msg_types:
+            self.shell_handlers[msg_type] = getattr(self.comm_manager, msg_type)
+        self.shell = SolaraInteractiveShell()
+        self.shell.display_pub.session = self.session
+        self.shell.display_pub.pub_socket = self.iopub_socket
+
+    async def _flush_control_queue(self):
+        pass
+
+    # these don't work from non-main thread, and we do not care about them I think
+    # TODO: it seems that if post_handler_hook is not override, the flask reload tests fails
+    # for unknown reason
+    def pre_handler_hook(self, *args):
+        pass
+
+    def post_handler_hook(self, *args):
+        pass
+
+    def set_parent(self, ident, parent, channel="shell"):
+        """Overridden from parent to tell the display hook and output streams
+        about the parent message.
+        """
+        if ipykernel_major < 6:
+            # the channel argument was added in 6.0
+            super().set_parent(ident, parent)
+        else:
+            super().set_parent(ident, parent, channel)
+        if channel == "shell":
+            self.shell.set_parent(parent)

solara/server/kernel_context.py

@@ -0,0 +1,357 @@
+import asyncio
+import sys
+
+try:
+    import contextvars
+except ModuleNotFoundError:
+    contextvars = None  # type: ignore
+
+import dataclasses
+import enum
+import logging
+import os
+import pickle
+import threading
+import time
+import typing
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, cast
+
+import ipywidgets as widgets
+import reacton
+from ipywidgets import DOMWidget, Widget
+
+import solara.server.settings
+import solara.util
+
+from . import kernel, kernel_context, websocket
+from .. import lifecycle
+from .kernel import Kernel, WebsocketStreamWrapper
+
+WebSocket = Any
+logger = logging.getLogger("solara.server.app")
+
+
+class Local(threading.local):
+    kernel_context_stack: Optional[List[Optional["kernel_context.VirtualKernelContext"]]] = None
+
+
+local = Local()
+
+
+class PageStatus(enum.Enum):
+    CONNECTED = "connected"
+    DISCONNECTED = "disconnected"
+    CLOSED = "closed"
+
+
+@dataclasses.dataclass
+class VirtualKernelContext:
+    id: str
+    kernel: kernel.Kernel
+    # we keep track of the session id to prevent kernel hijacking
+    # to 'steal' a kernel, one would need to know the session id
+    # *and* the kernel id
+    session_id: str
+    control_sockets: List[WebSocket] = dataclasses.field(default_factory=list)
+    # this is the 'private' version of the normally global ipywidgets.Widgets.widget dict
+    # see patch.py
+    widgets: Dict[str, Widget] = dataclasses.field(default_factory=dict)
+    # same, for ipyvue templates
+    # see patch.py
+    templates: Dict[str, Widget] = dataclasses.field(default_factory=dict)
+    user_dicts: Dict[str, Dict] = dataclasses.field(default_factory=dict)
+    # anything we need to attach to the context
+    # e.g. for a react app the render context, so that we can store/restore the state
+    app_object: Optional[Any] = None
+    reload: Callable = lambda: None  # noqa: E731
+    state: Any = None
+    container: Optional[DOMWidget] = None
+    # we track which pages are connected to implement kernel culling
+    page_status: Dict[str, PageStatus] = dataclasses.field(default_factory=dict)
+    # only used for testing
+    _last_kernel_cull_task: "Optional[asyncio.Future[None]]" = None
+    closed_event: threading.Event = dataclasses.field(default_factory=threading.Event)
+    _on_close_callbacks: List[Callable[[], None]] = dataclasses.field(default_factory=list)
+
+    def __post_init__(self):
+        with self:
+            for f, *_ in lifecycle._on_kernel_start_callbacks:
+                cleanup = f()
+                if cleanup:
+                    self.on_close(cleanup)
+
+    def restart(self):
+        # should we do this, or maybe close the context and create a new one?
+        with self:
+            for f in reversed(self._on_close_callbacks):
+                f()
+            self._on_close_callbacks.clear()
+            self.__post_init__()
+
+    def display(self, *args):
+        print(args)  # noqa
+
+    def on_close(self, f: Callable[[], None]):
+        self._on_close_callbacks.append(f)
+
+    def __enter__(self):
+        if local.kernel_context_stack is None:
+            local.kernel_context_stack = []
+        key = get_current_thread_key()
+        local.kernel_context_stack.append(current_context.get(key, None))
+        current_context[key] = self
+
+    def __exit__(self, *args):
+        key = get_current_thread_key()
+        assert local.kernel_context_stack is not None
+        current_context[key] = local.kernel_context_stack.pop()
+
+    def close(self):
+        if self.closed_event.is_set():
+            logger.error("Tried to close a kernel context that is already closed: %s", self.id)
+            return
+        logger.info("Shut down virtual kernel: %s", self.id)
+        with self:
+            for f in reversed(self._on_close_callbacks):
+                f()
+        with self:
+            if self.app_object is not None:
+                if isinstance(self.app_object, reacton.core._RenderContext):
+                    try:
+                        self.app_object.close()
+                    except Exception as e:
+                        logger.exception("Could not close render context: %s", e)
+                        # we want to continue, so we at least close all widgets
+            widgets.Widget.close_all()
+            # what if we reference each other
+            # import gc
+            # gc.collect()
+            self.kernel.session.close()
+        if self.id in contexts:
+            del contexts[self.id]
+        self.closed_event.set()
+
+    def _state_reset(self):
+        state_directory = Path(".") / "states"
+        state_directory.mkdir(exist_ok=True)
+        path = state_directory / f"{self.id}.pickle"
+        path = path.absolute()
+        try:
+            path.unlink()
+        except:  # noqa
+            pass
+        del contexts[self.id]
+        key = get_current_thread_key()
+        del current_context[key]
+
+    def state_save(self, state_directory: os.PathLike):
+        path = Path(state_directory) / f"{self.id}.pickle"
+        render_context = self.app_object
+        if render_context is not None:
+            render_context = cast(reacton.core._RenderContext, render_context)
+            state = render_context.state_get()
+            with path.open("wb") as f:
+                logger.debug("State: %r", state)
+                pickle.dump(state, f)
+
+    def page_connect(self, page_id: str):
+        logger.info("Connect page %s for kernel %s", page_id, self.id)
+        assert self.page_status.get(page_id) != PageStatus.CLOSED, "cannot connect with the same page_id after a close"
+        self.page_status[page_id] = PageStatus.CONNECTED
+        if self._last_kernel_cull_task:
+            self._last_kernel_cull_task.cancel()
+
+    def page_disconnect(self, page_id: str) -> "asyncio.Future[None]":
+        """Signal that a page has disconnected, and schedule a kernel cull if needed.
+
+        During the kernel reconnect window, we will keep the kernel alive, even if all pages have disconnected.
+
+        Returns a future that is set when the kernel cull is done.
+        The scheduled kernel cull can be cancelled when a new page connects, a new disconnect is scheduled,
+        or a page if explicitly closed.
+        """
+        logger.info("Disconnect page %s for kernel %s", page_id, self.id)
+        future: "asyncio.Future[None]" = asyncio.Future()
+        self.page_status[page_id] = PageStatus.DISCONNECTED
+        current_event_loop = asyncio.get_event_loop()
+
+        async def kernel_cull():
+            try:
+                cull_timeout_sleep_seconds = solara.util.parse_timedelta(solara.server.settings.kernel.cull_timeout)
+                logger.info("Scheduling kernel cull, will wait for max %s before shutting down the virtual kernel %s", cull_timeout_sleep_seconds, self.id)
+                await asyncio.sleep(cull_timeout_sleep_seconds)
+                has_connected_pages = PageStatus.CONNECTED in self.page_status.values()
+                if has_connected_pages:
+                    logger.info("We have (re)connected pages, keeping the virtual kernel %s alive", self.id)
+                else:
+                    logger.info("No connected pages, and timeout reached, shutting down virtual kernel %s", self.id)
+                    self.close()
+                current_event_loop.call_soon_threadsafe(future.set_result, None)
+            except asyncio.CancelledError:
+                if sys.version_info >= (3, 9):
+                    current_event_loop.call_soon_threadsafe(future.cancel, "cancelled because a new cull task was scheduled")
+                else:
+                    current_event_loop.call_soon_threadsafe(future.cancel)
+                raise
+
+        has_connected_pages = PageStatus.CONNECTED in self.page_status.values()
+        if not has_connected_pages:
+            # when we have no connected pages, we will schedule a kernel cull
+            if self._last_kernel_cull_task:
+                self._last_kernel_cull_task.cancel()
+
+            async def create_task():
+                task = asyncio.create_task(kernel_cull())
+                # create a reference to the task so we can cancel it later
+                self._last_kernel_cull_task = task
+                await task
+
+            asyncio.run_coroutine_threadsafe(create_task(), keep_alive_event_loop)
+        else:
+            future.set_result(None)
+        return future
+
+    def page_close(self, page_id: str):
+        """Signal that a page has closed, and close the context if needed.
+
+        Closing the browser tab or a page navigation means an explicit close, which is
+        different from a websocket/page disconnect, which we might want to recover from.
+
+        """
+        self.page_status[page_id] = PageStatus.CLOSED
+        logger.info("Close page %s for kernel %s", page_id, self.id)
+        has_connected_pages = PageStatus.CONNECTED in self.page_status.values()
+        has_disconnected_pages = PageStatus.DISCONNECTED in self.page_status.values()
+        if not (has_connected_pages or has_disconnected_pages):
+            logger.info("No connected or disconnected pages, shutting down virtual kernel %s", self.id)
+            if self._last_kernel_cull_task:
+                self._last_kernel_cull_task.cancel()
+            self.close()
+
+
+try:
+    # Normal Python
+    keep_alive_event_loop = asyncio.new_event_loop()
+
+    def _run():
+        asyncio.set_event_loop(keep_alive_event_loop)
+        try:
+            keep_alive_event_loop.run_forever()
+        except Exception:
+            logger.exception("Error in keep alive event loop")
+            raise
+
+    threading.Thread(target=_run, daemon=True).start()
+except RuntimeError:
+    # Emscripten/pyodide/lite
+    keep_alive_event_loop = asyncio.get_event_loop()
+
+contexts: Dict[str, VirtualKernelContext] = {}
+# maps from thread key to VirtualKernelContext, if VirtualKernelContext is None, it exists, but is not set as current
+current_context: Dict[str, Optional[VirtualKernelContext]] = {}
+
+
+def create_dummy_context():
+    from . import kernel
+
+    kernel_context = VirtualKernelContext(
+        id="dummy",
+        session_id="dummy",
+        kernel=kernel.Kernel(),
+    )
+    return kernel_context
+
+
+if contextvars is not None:
+    if typing.TYPE_CHECKING:
+        async_context_id = contextvars.ContextVar[str]("async_context_id")
+    else:
+        async_context_id = contextvars.ContextVar("async_context_id")
+    async_context_id.set("default")
+else:
+    async_context_id = None
+
+
+def get_current_thread_key() -> str:
+    if not solara.server.settings.kernel.threaded:
+        if async_context_id is not None:
+            try:
+                key = async_context_id.get()
+            except LookupError:
+                raise RuntimeError("no kernel context set")
+        else:
+            raise RuntimeError("No threading support, and no contextvars support (Python 3.6 is not supported for this)")
+    else:
+        thread = threading.current_thread()
+        key = get_thread_key(thread)
+    return key
+
+
+def get_thread_key(thread: threading.Thread) -> str:
+    if not solara.server.settings.kernel.threaded:
+        if async_context_id is not None:
+            return async_context_id.get()
+    thread_key = thread._name + str(thread._ident)  # type: ignore
+    return thread_key
+
+
+def set_context_for_thread(context: VirtualKernelContext, thread: threading.Thread):
+    key = get_thread_key(thread)
+    current_context[key] = context
+
+
+def has_current_context() -> bool:
+    thread_key = get_current_thread_key()
+    return (thread_key in current_context) and (current_context[thread_key] is not None)
+
+
+def get_current_context() -> VirtualKernelContext:
+    thread_key = get_current_thread_key()
+    if thread_key not in current_context:
+        raise RuntimeError(
+            f"Tried to get the current context for thread {thread_key}, but no known context found. This might be a bug in Solara. "
+            f"(known contexts: {list(current_context.keys())}"
+        )
+    context = current_context[thread_key]
+    if context is None:
+        raise RuntimeError(
+            f"Tried to get the current context for thread {thread_key!r}, although the context is know, it was not set for this thread. "
+            + "This might be a bug in Solara."
+        )
+    return context
+
+
+def set_current_context(context: Optional[VirtualKernelContext]):
+    thread_key = get_current_thread_key()
+    current_context[thread_key] = context
+
+
+def initialize_virtual_kernel(session_id: str, kernel_id: str, websocket: websocket.WebsocketWrapper):
+    from solara.server import app as appmodule
+
+    if kernel_id in contexts:
+        logger.info("reusing virtual kernel: %s", kernel_id)
+        context = contexts[kernel_id]
+        if context.session_id != session_id:
+            logger.critical("Session id mismatch when reusing kernel (hack attempt?): %s != %s", context.session_id, session_id)
+            websocket.send_text("Session id mismatch when reusing kernel (hack attempt?)")
+            # to avoid very fast reconnects (we are in a thread anyway)
+            time.sleep(0.5)
+            raise ValueError("Session id mismatch")
+        kernel = context.kernel
+    else:
+        kernel = Kernel()
+        logger.info("new virtual kernel: %s", kernel_id)
+        context = contexts[kernel_id] = VirtualKernelContext(id=kernel_id, session_id=session_id, kernel=kernel, control_sockets=[], widgets={}, templates={})
+
+        with context:
+            widgets.register_comm_target(kernel)
+            appmodule.register_solara_comm_target(kernel)
+    with context:
+        assert has_current_context()
+        assert kernel is Kernel.instance()
+        kernel.shell_stream = WebsocketStreamWrapper(websocket, "shell")
+        kernel.control_stream = WebsocketStreamWrapper(websocket, "control")
+        kernel.session.websockets.add(websocket)
+    return context