ENH: enhanced handling of IPython kernel parameters

Author: dmgav

bluesky_queueserver/manager/worker.py

@@ -1808,30 +1808,82 @@ def find_kernel_ip(ip_str):
                     ip = ip_str
                 return ip
 
-            shell_port = None ##60000
-            iopub_port = None ##60001
-            stdin_port = None ##60002
-            hb_port = None ##60003
-            control_port = None ##60004
-            # connection_file = "test.json"
+            # The following parameters should be properly added to CLI and config file parameters
+            connection_file = os.environ.get("QSERVER_IPYTHON_KERNEL_CONNECTION_FILE", None)
+            connection_dir = os.environ.get("QSERVER_IPYTHON_KERNEL_CONNECTION_DIR", None)
+            shell_port = os.environ.get("QSERVER_IPYTHON_KERNEL_SHELL_PORT", None)
+            iopub_port = os.environ.get("QSERVER_IPYTHON_KERNEL_IOPUB_PORT", None)
+            stdin_port = os.environ.get("QSERVER_IPYTHON_KERNEL_STDIN_PORT", None)
+            hb_port = os.environ.get("QSERVER_IPYTHON_KERNEL_HB_PORT", None)
+            control_port = os.environ.get("QSERVER_IPYTHON_KERNEL_CONTROL_PORT", None)
+
+            def _convert_to_int(x):
+                return int(x) if x else None
+
+            shell_port = _convert_to_int(shell_port)
+            iopub_port = _convert_to_int(iopub_port)
+            stdin_port = _convert_to_int(stdin_port)
+            hb_port = _convert_to_int(hb_port)
+            control_port = _convert_to_int(control_port)
 
-            # self._ip_kernel_app.connection_file = connection_file
+            kernel_ip = self._config_dict["ipython_kernel_ip"]
+            kernel_ip = find_kernel_ip(kernel_ip)
+
+            use_connection_file = bool(connection_file)
+
+            if connection_dir:
+                self._ip_kernel_app.connection_dir = connection_dir
+
+            if connection_file:
+                self._ip_kernel_app.connection_file = connection_file
+                abs_cf_name = self._ip_kernel_app.abs_connection_file
+
+                # Check the connection file. Delete the existing connection file if
+                #   any of the parameters are not matching the new parameters or if
+                #   the file is corrupt.
+                if os.path.isfile(abs_cf_name):
+                    try:
+                        with open(abs_cf_name, "r") as f:
+                            cn_info = json.load(f)
+
+                        def _check_value(value, key):
+                            if value:
+                                if key not in cn_info:
+                                    raise Exception(f"Key {key!r} is not found in the connection file")
+                                if cn_info[key] != value:
+                                    raise Exception(
+                                        f"Old value {cn_info[key]=} is does not match the new value {value!r}"
+                                    )
+
+                        _check_value(shell_port, "shell_port")
+                        _check_value(iopub_port, "iopub_port")
+                        _check_value(stdin_port, "stdin_port")
+                        _check_value(hb_port, "hb_port")
+                        _check_value(control_port, "control_port")
+                        _check_value(kernel_ip, "ip")
+
+                    except Exception as ex:
+                        logger.error(
+                            f"Connection file {abs_cf_name!r} is can't be loaded or out of date. "
+                            f"A new connection file will be generated. ({ex})"
+                        )
+                        use_connection_file = False
+                        os.remove(abs_cf_name)
+                else:
+                    use_connection_file = False
 
             logger.info("Generating random port numbers for IPython kernel ...")
-            kernel_ip = self._config_dict["ipython_kernel_ip"]
             try:
-                kernel_ip = find_kernel_ip(kernel_ip)
-                self._ip_kernel_app.ip = kernel_ip
-                self._ip_kernel_app.shell_port = generate_random_port(kernel_ip)
-                self._ip_kernel_app.iopub_port = generate_random_port(kernel_ip)
-                self._ip_kernel_app.stdin_port = generate_random_port(kernel_ip)
-                self._ip_kernel_app.hb_port = generate_random_port(kernel_ip)
-                self._ip_kernel_app.control_port = control_port or generate_random_port(kernel_ip)
-                # self._ip_kernel_app.shell_port = shell_port or generate_random_port(kernel_ip)
-                # self._ip_kernel_app.iopub_port = iopub_port or generate_random_port(kernel_ip)
-                # self._ip_kernel_app.stdin_port = stdin_port or generate_random_port(kernel_ip)
-                # self._ip_kernel_app.hb_port = hb_port or generate_random_port(kernel_ip)
-                # self._ip_kernel_app.control_port = control_port or generate_random_port(kernel_ip)
+                if use_connection_file:
+                    logger.info(f"Loading connection file {self._ip_kernel_app.abs_connection_file}")
+                    self._ip_kernel_app.load_connection_file(self._ip_kernel_app.abs_connection_file)
+                else:
+                    self._ip_kernel_app.ip = kernel_ip
+                    self._ip_kernel_app.shell_port = shell_port or generate_random_port(kernel_ip)
+                    self._ip_kernel_app.iopub_port = iopub_port or generate_random_port(kernel_ip)
+                    self._ip_kernel_app.stdin_port = stdin_port or generate_random_port(kernel_ip)
+                    self._ip_kernel_app.hb_port = hb_port or generate_random_port(kernel_ip)
+                    self._ip_kernel_app.control_port = control_port or generate_random_port(kernel_ip)
                 self._ip_connect_info = self._ip_kernel_app.get_connection_info()
             except Exception as ex:
                 self._success_startup = False