Merge pull request #207 from Keck-DataReductionPipelines/develop

Develop

Author: rmcgurk

kcwidrp/configs/kcwi.cfg

@@ -23,6 +23,17 @@ enable_bokeh = True
 plot_level = 1
 # How long to pause between plots at plot_level = 1 (seconds)
 plot_pause = 1
+# Use the Firefox-compatible plot export path. The simple export (False) works
+# for most users; set to True only on systems where Bokeh's default browser
+# detection fails (e.g. Ubuntu 24.04 with snap-confined Firefox).
+# Stop the pipeline if the Bokeh server fails to start. When False, the
+# pipeline continues without interactive plots and logs a warning instead.
+terminate_on_failed_bokeh_start = False
+plot_firefox_compat = False
+# Pre-warm the Firefox WebDriver at pipeline startup so the first save_plot
+# call does not pay the browser startup cost. Only relevant when
+# plot_firefox_compat = True.
+plot_prewarm_firefox = False
 # Wavelength fitting verbosity: 1 - normal, 2 - verbose
 verbose = 1
 # Bokeh plot dimensions in pixels

kcwidrp/core/bokeh_plotting.py

@@ -20,6 +20,9 @@ def bokeh_plot(plot, session):
 
     # NOT TESTED YET
 
+    if session is None:
+        return
+
     new_figure = session.document.select_one(selector=dict(type=Figure))
     layout = session.document.select_one(selector=dict(type=Column))
 

kcwidrp/core/kcwi_plotting.py

@@ -1,11 +1,72 @@
 import numpy as np
 from bokeh.io import export_png
 
+import atexit
 import os
+import shutil
 import logging
 
+from selenium import webdriver
+
 logger = logging.getLogger('KCWI')
 
+_firefox_driver = None
+
+
+def configure_plot_driver(firefox_compat=False, prewarm=False):
+    """Optionally pre-warm the Firefox WebDriver at pipeline startup.
+
+    Called by StartBokeh when both plot_firefox_compat and plot_prewarm_firefox
+    are True in kcwi.cfg.  Once the driver is created here, save_plot will
+    automatically use the Firefox-compatible export path for all subsequent
+    calls.
+    """
+    if firefox_compat and prewarm:
+        logger.info("Pre-warming Firefox WebDriver for plot export")
+        _get_driver()
+
+
+def _get_driver():
+    """Return a cached Firefox WebDriver, creating it on first call."""
+    global _firefox_driver
+    if _firefox_driver is not None:
+        try:
+            # Navigate to about:blank to cancel any pending navigation left by
+            # Bokeh's reset step (it navigates to http://0.0.0.1/ after each
+            # screenshot; on some networks that address times out rather than
+            # refusing immediately, blocking the next export call).
+            _firefox_driver.get("about:blank")
+            return _firefox_driver
+        except Exception:
+            _firefox_driver = None
+
+    from selenium.webdriver.firefox.service import Service as FirefoxService
+    options = webdriver.FirefoxOptions()
+    options.add_argument("--headless")
+    options.add_argument("--no-sandbox")
+    options.add_argument("--disable-dev-shm-usage")
+    options.add_argument("--hide-scrollbars")
+    options.add_argument("--force-device-scale-factor=1")
+    options.add_argument("--force-color-profile=srgb")
+    firefox_bin = _find_firefox_binary()
+    if firefox_bin:
+        options.binary_location = firefox_bin
+    geckodriver = _find_geckodriver()
+    serv = FirefoxService(executable_path=geckodriver) if geckodriver else FirefoxService()
+    _firefox_driver = webdriver.Firefox(options=options, service=serv)
+    atexit.register(_close_driver)
+    return _firefox_driver
+
+
+def _close_driver():
+    global _firefox_driver
+    if _firefox_driver is not None:
+        try:
+            _firefox_driver.quit()
+        except Exception:
+            pass
+        _firefox_driver = None
+
 
 def get_plot_lims(data, padding=0.05, clip=True):
     """Get plot limits using data range plus padding fraction"""
@@ -33,12 +94,78 @@ def set_plot_lims(fig, xlim=None, ylim=None):
         fig.y_range.start = ylim[0]
         fig.y_range.end = ylim[1]
 
+def _find_firefox_binary():
+    """Return path to the real Firefox binary, handling snap wrapper installs."""
+    # On Ubuntu 24.04, /usr/bin/firefox is a shell script wrapper for the snap.
+    # Selenium requires the actual ELF binary, not a wrapper script.
+    candidates = [
+        '/snap/firefox/current/usr/lib/firefox/firefox',
+        '/usr/lib/firefox/firefox',
+    ]
+    for path in candidates:
+        if os.path.isfile(path) and os.access(path, os.X_OK):
+            return path
+    # Fall back to whatever is on PATH (works on most non-snap installs)
+    return shutil.which('firefox')
+
+
+def _find_geckodriver():
+    """Return path to geckodriver, searching common locations."""
+    candidates = [
+        '/snap/bin/geckodriver',
+        '/usr/bin/geckodriver',
+        '/usr/local/bin/geckodriver',
+    ]
+    for path in candidates:
+        if os.path.isfile(path) and os.access(path, os.X_OK):
+            return path
+    return shutil.which('geckodriver')
+
 
 def save_plot(fig, filename=None):
+    """Save a Bokeh figure to a PNG file.
+
+    Dispatches to the Firefox-compatible path or the simple path depending on
+    whether configure_plot_driver() was called with firefox_compat=True.
+    """
+    if _firefox_driver is not None:
+        _save_plot_firefox(fig, filename)
+    else:
+        _save_plot_simple(fig, filename)
+
+
+def _save_plot_firefox(fig, filename=None):
+    """Firefox-compatible export for systems with snap-confined Firefox."""
     if filename is None:
         fnam = os.path.join('plots', 'kcwi_drp_plot.png')
     else:
         fnam = os.path.join('plots', filename)
-    export_png(fig, filename=fnam)
 
+    # Resolve to absolute path before chdir so the PNG lands in the right place.
+    fnam = os.path.abspath(fnam)
+    os.makedirs(os.path.dirname(fnam), exist_ok=True)
+
+    driver = _get_driver()
+    # Snap-confined Firefox on Ubuntu 24.04 uses a private /tmp mount namespace,
+    # so file:///tmp/... URLs are invisible to it.  $HOME is accessible via the
+    # snap 'home' interface.  Bokeh writes its scratch HTML relative to the
+    # process CWD, so chdir to $HOME to land the temp file somewhere Firefox
+    # can actually reach.  fnam is already absolute so the PNG output goes to
+    # the right place regardless.
+    old_cwd = os.getcwd()
+    try:
+        os.chdir(os.path.expanduser('~'))
+        export_png(fig, filename=fnam, webdriver=driver)
+    finally:
+        os.chdir(old_cwd)
+    logger.info(">>> Saving to %s" % fnam)
+
+
+def _save_plot_simple(fig, filename=None):
+    """Simple export using Bokeh's built-in browser detection."""
+    if filename is None:
+        fnam = os.path.join('plots', 'kcwi_drp_plot.png')
+    else:
+        fnam = os.path.join('plots', filename)
+    export_png(fig, filename=fnam)
     logger.info(">>> Saving to %s" % fnam)

kcwidrp/primitives/StartBokeh.py

@@ -4,6 +4,8 @@
 from bokeh.plotting.figure import figure
 from bokeh.layouts import column
 
+from kcwidrp.core.kcwi_plotting import configure_plot_driver
+
 
 class StartBokeh(BasePrimitive):
     """
@@ -23,8 +25,18 @@ def __init__(self, action, context):
 
     def _perform(self):
 
-        # session = pull_session(session_id='kcwi', url='http://localhost:5006')
-        session = pull_session()
+        self.context.bokeh_session = None
+        try:
+            # session = pull_session(session_id='kcwi', url='http://localhost:5006')
+            session = pull_session()
+        except Exception as e:
+            if getattr(self.config.instrument, 'terminate_on_failed_bokeh_start', False):
+                self.logger.error("Could not connect to Bokeh server: %s", e)
+                raise RuntimeError("Bokeh server failed to start and "
+                                   "terminate_on_failed_bokeh_start is True") from e
+            self.logger.warning("Could not connect to Bokeh server, "
+                                "interactive plots disabled: %s", e)
+            return self.action.args
         self.logger.info("Enabling BOKEH plots")
         p = figure()
         c = column(children=[p])
@@ -33,4 +45,9 @@ def _perform(self):
         self.context.bokeh_session = session
         session.show(c)
 
+        firefox_compat = getattr(self.config.instrument, 'plot_firefox_compat', False)
+        prewarm = getattr(self.config.instrument, 'plot_prewarm_firefox', False)
+        if firefox_compat and prewarm:
+            configure_plot_driver(firefox_compat=firefox_compat, prewarm=prewarm)
+
         return True

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kcwidrp"
-version = "1.3.0"
+version = "1.3.1"
 description = "KCWI Data Reduction Pipeline"
 readme = "README.rst"
 license = "BSD-3-Clause"