jupyter kernel ports are now determined in pyscript and made available to hass_pyscript_kernel.py via a state variable

Author: craigbarratt

custom_components/pyscript/__init__.py

@@ -1,6 +1,7 @@
 """Component to allow running Python scripts."""
 
 import glob
+import json
 import logging
 import os
 
@@ -114,6 +115,12 @@ async def jupyter_kernel_start(call):
         handler_func.install_ast_funcs(ast_ctx)
         kernel = Kernel(call.data, ast_ctx, global_ctx_name, global_ctx_mgr)
         await kernel.session_start()
+        hass.states.async_set(call.data["state_var"], json.dumps(kernel.get_ports()))
+
+        def state_var_remove():
+            hass.states.async_remove(call.data["state_var"])
+
+        kernel.set_session_cleanup_callback(state_var_remove)
 
     hass.services.async_register(
         DOMAIN, SERVICE_JUPYTER_KERNEL_START, jupyter_kernel_start
@@ -123,12 +130,10 @@ async def state_changed(event):
         var_name = event.data["entity_id"]
         # attr = event.data["new_state"].attributes
         if "new_state" not in event.data or event.data["new_state"] is None:
-            _LOGGER.debug(
-                "state_changed: missing new_state in event.data=%s; ignoring",
-                event.data,
-            )
-            return
-        new_val = event.data["new_state"].state
+            # state variable has been deleted
+            new_val = None
+        else:
+            new_val = event.data["new_state"].state
         old_val = event.data["old_state"].state if event.data["old_state"] else None
         new_vars = {var_name: new_val, f"{var_name}.old": old_val}
         func_args = {

custom_components/pyscript/jupyter_kernel.py

@@ -14,6 +14,7 @@
 import json
 import logging
 import logging.handlers
+import random
 import re
 from struct import pack, unpack
 import traceback
@@ -36,14 +37,6 @@ def str_to_bytes(string):
     """Encode a string in bytes."""
     return string.encode('ascii')
 
-def bind(socket, connection, port):
-    """Bind a socket."""
-    if port <= 0:
-        return socket.bind_to_random_port(connection)
-    # _LOGGER.debug(f"binding to %s:%s" % (connection, port))
-    socket.bind("%s:%s" % (connection, port))
-    return port
-
 class KernelBufferingHandler(logging.handlers.BufferingHandler):
     """Memory-based handler for logging; send via stdout queue."""
     def __init__(self, housekeep_q):
@@ -61,7 +54,6 @@ def shouldFlush(self, record):
 
 
 ################################################################
-#
 class ZmqSocket:
     """Defines a minimal implementation of a small subset of ZMQ,
 allowing pyscript to work with Jupyter without the real zmq
@@ -193,7 +185,6 @@ def __init__(self, config, ast_ctx, global_ctx_name, global_ctx_mgr):
         self.global_ctx_mgr = global_ctx_mgr
         self.ast_ctx = ast_ctx
 
-        self.connection = self.config["transport"] + "://" + self.config["ip"]
         self.secure_key = str_to_bytes(self.config["key"])
         self.signature_schemes = {"hmac-sha256": hashlib.sha256}
         self.auth = hmac.HMAC(
@@ -209,17 +200,30 @@ def __init__(self, config, ast_ctx, global_ctx_name, global_ctx_mgr):
         self.stdin_server = None
         self.shell_server = None
 
+        self.heartbeat_port = None
+        self.iopub_port = None
+        self.control_port = None
+        self.stdin_port = None
+        self.shell_port = None
+        self.avail_port = random.randrange(40000, 50000)
+
         # there can be multiple iopub subscribers, with corresponding tasks
         self.iopub_socket = set()
 
         self.tasks = {}
         self.task_cnt = 0
         self.task_cnt_max = 0
 
+        self.session_cleanup_callback = None
+
         self.housekeep_q = asyncio.Queue(0)
 
         self.parent_header = None
 
+        #
+        # we create a logging handler so that output from the log functions
+        # gets delivered back to Jupyter as stdout
+        #
         self.console = KernelBufferingHandler(self.housekeep_q)
         self.console.setLevel(logging.DEBUG)
         # set a format which is just the message
@@ -514,13 +518,10 @@ async def shell_handler(self, shell_socket, msg):
         else:
             _LOGGER.error("unknown msg_type: %s", msg['header']["msg_type"])
 
-
-    ##########################################
-    # Control:
     async def control_listen(self, reader, writer):
         """Task that listens to control messages."""
         try:
-            # _LOGGER.debug("control_listen connected")
+            _LOGGER.debug("control_listen connected")
             await self.housekeep_q.put(["register", "control", current_task()])
             control_socket = ZmqSocket(reader, writer, "ROUTER")
             await control_socket.handshake()
@@ -533,41 +534,37 @@ async def control_listen(self, reader, writer):
         except asyncio.CancelledError:  # pylint: disable=try-except-raise
             raise
         except EOFError:
-            # _LOGGER.debug("control_listen got eof")
+            _LOGGER.debug("control_listen got eof")
             await self.housekeep_q.put(["unregister", "control", current_task()])
             control_socket.close()
         except Exception as err:  # pylint: disable=broad-except
             _LOGGER.error("control_listen exception %s", err)
             await self.housekeep_q.put(["shutdown"])
 
-    ##########################################
-    # Stdin:
     async def stdin_listen(self, reader, writer):
         """Task that listens to stdin messages."""
         try:
-            # _LOGGER.debug("stdin_listen connected")
+            _LOGGER.debug("stdin_listen connected")
             await self.housekeep_q.put(["register", "stdin", current_task()])
             stdin_socket = ZmqSocket(reader, writer, "ROUTER")
             await stdin_socket.handshake()
             while 1:
                 _ = await stdin_socket.recv_multipart()
-                # _LOGGER.debug("stdin_listen received %s", raw_msg)
+                # _LOGGER.debug("stdin_listen received %s", _)
         except asyncio.CancelledError:  # pylint: disable=try-except-raise
             raise
         except EOFError:
-            # _LOGGER.debug("stdin_listen got eof")
+            _LOGGER.debug("stdin_listen got eof")
             await self.housekeep_q.put(["unregister", "stdin", current_task()])
             stdin_socket.close()
         except Exception:  # pylint: disable=broad-except
             _LOGGER.error("stdin_listen exception %s", traceback.format_exc(-1))
             await self.housekeep_q.put(["shutdown"])
 
-    ##########################################
-    # Shell:
     async def shell_listen(self, reader, writer):
         """Task that listens to shell messages."""
         try:
-            # _LOGGER.debug("shell_listen connected")
+            _LOGGER.debug("shell_listen connected")
             await self.housekeep_q.put(["register", "shell", current_task()])
             shell_socket = ZmqSocket(reader, writer, "ROUTER")
             await shell_socket.handshake()
@@ -578,62 +575,56 @@ async def shell_listen(self, reader, writer):
             shell_socket.close()
             raise
         except EOFError:
-            # _LOGGER.debug("shell_listen got eof")
+            _LOGGER.debug("shell_listen got eof")
             await self.housekeep_q.put(["unregister", "shell", current_task()])
             shell_socket.close()
         except Exception:  # pylint: disable=broad-except
             _LOGGER.error("shell_listen exception %s", traceback.format_exc(-1))
             await self.housekeep_q.put(["shutdown"])
 
-    ##########################################
-    # Heartbeat:
     async def heartbeat_listen(self, reader, writer):
         """Task that listens and responds to heart beat messages."""
         try:
-            # _LOGGER.debug("heartbeat_listen connected")
+            _LOGGER.debug("heartbeat_listen connected")
             await self.housekeep_q.put(["register", "heartbeat", current_task()])
             heartbeat_socket = ZmqSocket(reader, writer, "REP")
             await heartbeat_socket.handshake()
             while 1:
                 msg = await heartbeat_socket.recv()
-                # _LOGGER.debug(f"heartbeat_listen: got {msg}")
+                # _LOGGER.debug("heartbeat_listen: got %s", msg)
                 await heartbeat_socket.send(msg)
         except asyncio.CancelledError:  # pylint: disable=try-except-raise
             raise
         except EOFError:
-            # _LOGGER.debug("heartbeat_listen got eof")
+            _LOGGER.debug("heartbeat_listen got eof")
             await self.housekeep_q.put(["unregister", "heartbeat", current_task()])
             heartbeat_socket.close()
         except Exception:  # pylint: disable=broad-except
             _LOGGER.error("heartbeat_listen exception: %s", traceback.format_exc(-1))
             await self.housekeep_q.put(["shutdown"])
 
-    ##########################################
-    # IOPub/Sub:
     async def iopub_listen(self, reader, writer):
         """Task that listens to iopub messages."""
         try:
-            # _LOGGER.debug("iopub_listen connected")
+            _LOGGER.debug("iopub_listen connected")
             await self.housekeep_q.put(["register", "iopub", current_task()])
             iopub_socket = ZmqSocket(reader, writer, "PUB")
             await iopub_socket.handshake()
             self.iopub_socket.add(iopub_socket)
             while 1:
                 _ = await iopub_socket.recv_multipart()
-                # _LOGGER.debug("iopub received %s", wire_msg)
+                # _LOGGER.debug("iopub received %s", _)
         except asyncio.CancelledError:  # pylint: disable=try-except-raise
             raise
         except EOFError:
             await self.housekeep_q.put(["unregister", "iopub", current_task()])
             iopub_socket.close()
             self.iopub_socket.discard(iopub_socket)
-            # _LOGGER.debug("iopub_listen got eof")
+            _LOGGER.debug("iopub_listen got eof")
         except Exception:  # pylint: disable=broad-except
             _LOGGER.error("iopub_listen exception %s", traceback.format_exc(-1))
             await self.housekeep_q.put(["shutdown"])
 
-    ##########################################
-    # Housekeeping
     async def housekeep_run(self):
         """Housekeeping, including closing servers after startup, and doing orderly shutdown."""
         while True:
@@ -652,6 +643,12 @@ async def housekeep_run(self):
                     self.tasks[msg[1]].add(msg[2])
                     self.task_cnt += 1
                     self.task_cnt_max = max(self.task_cnt_max, self.task_cnt)
+                    #
+                    # now a couple of things are connected, call the session_cleanup_callback
+                    #
+                    if self.task_cnt > 1 and self.session_cleanup_callback:
+                        self.session_cleanup_callback()
+                        self.session_cleanup_callback = None
                 elif msg[0] == "unregister":
                     if msg[1] in self.tasks:
                         self.tasks[msg[1]].discard(msg[2])
@@ -670,35 +667,89 @@ async def housekeep_run(self):
             except Exception:  # pylint: disable=broad-except
                 _LOGGER.error("housekeep task exception: %s", traceback.format_exc(-1))
 
+    async def startup_timeout(self):
+        """Shut down the session if nothing connects after 30 seconds."""
+        await self.housekeep_q.put(["register", "startup_timeout", current_task()])
+        await asyncio.sleep(30)
+        if self.task_cnt_max == 1:
+            #
+            # nothing started other than us, so shut down the session
+            #
+            _LOGGER.error("No connections to session %s; shutting down", self.global_ctx_name)
+            if self.session_cleanup_callback:
+                self.session_cleanup_callback()
+                self.session_cleanup_callback = None
+            await self.housekeep_q.put(["shutdown"])
+        await self.housekeep_q.put(["unregister", "startup_timeout", current_task()])
+
+    async def start_one_server(self, callback):
+        """Start a server by finding an available port."""
+        for _ in range(2048):
+            try:
+                server = await asyncio.start_server(callback, self.config["ip"], self.avail_port)
+                return server, self.avail_port
+            except OSError:
+                self.avail_port += 1
+        _LOGGER.error("unable to find an available port on host %s, last port %d", self.config["ip"], self.avail_port)
+        return None, None
+
+    def get_ports(self):
+        """Return a dict of the port numbers this kernel session is listening to."""
+        return {
+            "iopub_port": self.iopub_port,
+            "hb_port": self.heartbeat_port,
+            "control_port": self.control_port,
+            "stdin_port": self.stdin_port,
+            "shell_port": self.shell_port,
+        }
+
+    def set_session_cleanup_callback(self, callback):
+        """Set a cleanup callback which is called right after the session has started."""
+        self.session_cleanup_callback = callback
+
     async def session_start(self):
         """Start the kernel session."""
         self.ast_ctx.add_logger_handler(self.console)
         _LOGGER.info("Starting session %s", self.global_ctx_name)
 
         self.tasks["housekeep"] = {asyncio.create_task(self.housekeep_run())}
+        self.tasks["startup_timeout"] = {asyncio.create_task(self.startup_timeout())}
 
-        self.iopub_server = await asyncio.start_server(self.iopub_listen, self.config["ip"], self.config["iopub_port"])
-        self.heartbeat_server = await asyncio.start_server(self.heartbeat_listen, self.config["ip"], self.config["hb_port"])
-        self.control_server = await asyncio.start_server(self.control_listen, self.config["ip"], self.config["control_port"])
-        self.stdin_server = await asyncio.start_server(self.stdin_listen, self.config["ip"], self.config["stdin_port"])
-        self.shell_server = await asyncio.start_server(self.shell_listen, self.config["ip"], self.config["shell_port"])
+        self.iopub_server, self.iopub_port = await self.start_one_server(self.iopub_listen)
+        self.heartbeat_server, self.heartbeat_port = await self.start_one_server(self.heartbeat_listen)
+        self.control_server, self.control_port = await self.start_one_server(self.control_listen)
+        self.stdin_server, self.stdin_port = await self.start_one_server(self.stdin_listen)
+        self.shell_server, self.shell_port = await self.start_one_server(self.shell_listen)
 
         #
         # For debugging, can use the real ZMQ library instead on certain sockets; comment out
         # the corresponding asyncio.start_server() call above if you enable the ZMQ-based
-        # functions here.  The two most important ones are shown here.
+        # functions here.  You can then turn of verbosity level 4 (-vvvv) in hass_pyscript_kernel.py
+        # to see all the byte data in case you need to debug the simple ZMQ implementation here.
+        # The two most important zmq functions are shown below.
         #
         #  import zmq
         #  import zmq.asyncio
         #
+        #  def zmq_bind(socket, connection, port):
+        #      """Bind a socket."""
+        #      if port <= 0:
+        #          return socket.bind_to_random_port(connection)
+        #      # _LOGGER.debug(f"binding to %s:%s" % (connection, port))
+        #      socket.bind("%s:%s" % (connection, port))
+        #      return port
+        #
         #  zmq_ctx = zmq.asyncio.Context()
+        #
         #  ##########################################
         #  # Shell using real ZMQ for debugging:
         #  async def shell_listen_zmq():
         #      """Task that listens to shell messages using ZMQ."""
         #      try:
+        #          _LOGGER.debug("shell_listen_zmq connected")
+        #          connection = self.config["transport"] + "://" + self.config["ip"]
         #          shell_socket = zmq_ctx.socket(zmq.ROUTER)  # pylint: disable=no-member
-        #          self.config["shell_port"] = bind(shell_socket, self.connection, self.config["shell_port"])
+        #          self.shell_port = zmq_bind(shell_socket, connection, -1)
         #          _LOGGER.debug("shell_listen_zmq connected")
         #          while 1:
         #              msg = await shell_socket.recv_multipart()
@@ -716,8 +767,9 @@ async def session_start(self):
         #      """Task that listens to iopub messages using ZMQ."""
         #      try:
         #          _LOGGER.debug("iopub_listen_zmq connected")
+        #          connection = self.config["transport"] + "://" + self.config["ip"]
         #          iopub_socket = zmq_ctx.socket(zmq.PUB)  # pylint: disable=no-member
-        #          self.config["iopub_port"] = bind(self.iopub_socket, self.connection, self.config["iopub_port"])
+        #          self.iopub_port = zmq_bind(self.iopub_socket, connection, -1)
         #          self.iopub_socket.add(iopub_socket)
         #          while 1:
         #              wire_msg = await iopub_socket.recv_multipart()