benchlab-pytools 3.0.0__py3-none-any.whl

benchlab/wigidash/wigidash_manager.py

@@ -0,0 +1,342 @@
+import os
+import threading
+import time
+import types
+
+from benchlab.core.datasource_manager import DataSourceManager
+from benchlab.wigidash.benchlab_telemetry import (
+    TelemetryHistory, telemetry_step, TelemetryContext
+)
+from benchlab.wigidash.benchlab_utils import get_logger
+from benchlab.wigidash.wigidash_usb import scan_wigidash
+from benchlab.wigidash.wigidash_session import BenchlabWigiSession
+
+logger = get_logger("WigidashManager")
+
+
+class WigidashManager:
+    """
+    Manages multiple Wigidash sessions and their assigned Benchlab devices.
+    """
+
+    def __init__(self, vendor_id=0x28DA, product_id=0xEF01, datasource=None):
+        self.vendor_id = vendor_id
+        self.product_id = product_id
+        self.sessions = []
+        self.benchlab_devices = {}
+        self.telemetry_histories = {}
+        self.telemetry_contexts = {}
+        self.shutting_down = False
+        self.shutdown_event = threading.Event()
+        self.shutdown_barrier = None
+        self.datasource = datasource  # Already-connected DataSourceManager
+        self._telemetry_lock = threading.Lock()
+
+    # ----------------- BENCHLAB DEVICE MANAGEMENT ----------------- #
+    def get_available_benchlabs(self, log_info=True):
+        """Query the datasource for available Benchlab devices."""
+        devices = []
+        query_succeeded = False
+        try:
+            raw = self.datasource.list_devices()
+            if isinstance(raw, dict):
+                devices = [{"port": info.get("port", "?"), "uid": uid,
+                            "firmware": info.get("firmware", "?")}
+                           for uid, info in raw.items()]
+            else:
+                devices = list(raw)
+            query_succeeded = True
+        except Exception as e:
+            logger.warning(f"list_devices failed: {e}")
+
+        for d in devices:
+            port = d["port"]
+            uid = d["uid"]
+            fw = d.get("firmware", "?")
+            if port not in self.benchlab_devices:
+                self.benchlab_devices[port] = {
+                    "uid": uid, "firmware": fw, "in_use": False}
+            else:
+                self.benchlab_devices[port]["uid"] = uid
+                self.benchlab_devices[port]["firmware"] = fw
+
+        # Prune ports no longer reported by the data source, so a disconnected
+        # device doesn't linger forever with stale UID/firmware/in_use state.
+        if query_succeeded:
+            current_ports = {d["port"] for d in devices}
+            for stale_port in set(self.benchlab_devices) - current_ports:
+                del self.benchlab_devices[stale_port]
+                self.telemetry_contexts.pop(stale_port, None)
+
+        all_devices = [
+            {"port": port, "uid": info["uid"], "in_use": info["in_use"]}
+            for port, info in self.benchlab_devices.items()
+        ]
+
+        if log_info:
+            count = len(all_devices)
+            if count:
+                logger.info(
+                    f"{count} available Benchlab device{
+                        's' if count > 1 else ''}:")
+                for dev in all_devices:
+                    logger.info(f"  {dev['port']}: UID {dev['uid']}")
+            else:
+                logger.info("No Benchlab devices detected.")
+
+        return all_devices
+
+    # ----------------- TELEMETRY MANAGEMENT ----------------- #
+    def release_port(self, port):
+        if port in self.benchlab_devices:
+            self.benchlab_devices[port]["in_use"] = False
+
+    def start_telemetry(self, port, session: BenchlabWigiSession):
+        if port not in self.benchlab_devices:
+            logger.warning(f"Cannot start telemetry: unknown port {port}")
+            return
+
+        with self._telemetry_lock:
+            # ------------------------------------------------
+            # CASE 1: Telemetry already running for this port
+            # ------------------------------------------------
+            if port in self.telemetry_contexts:
+                ctx = self.telemetry_contexts[port]
+
+                logger.info(
+                    f"Selected device {port} is already in use. "
+                    "Using existing telemetry.")
+
+                session.ser = ctx.ser
+                session.device_info = ctx.device_info
+                session.uid = ctx.uid
+                session.telemetry_history = ctx.history
+                session.history = ctx.history
+                session.selected_port = port
+                session.telemetry_context = ctx
+
+                ctx.sessions.append(session)
+                return
+
+            # ------------------------------------------------
+            # CASE 2: First session → start telemetry
+            # ------------------------------------------------
+            self.benchlab_devices[port]["in_use"] = True
+            uid = self.benchlab_devices[port]["uid"]
+            history = self.telemetry_histories.setdefault(
+                uid, TelemetryHistory())
+
+            # Get device info from datasource
+            device_info = {}
+            try:
+                self.datasource.select_device(uid)
+                snap = self.datasource.snapshot()
+                device_info = snap.get("device_info") or snap.get(
+                    "all_devices", {}).get(uid, {})
+            except Exception as e:
+                logger.warning(f"Could not get device info for {uid}: {e}")
+
+            ctx = TelemetryContext(
+                port=port,
+                ser=None,
+                device_info=device_info,
+                uid=uid,
+                history=history,
+            )
+            ctx.sessions = [session]
+            self.telemetry_contexts[port] = ctx
+
+            session.ser = None
+            session.device_info = device_info
+            session.uid = uid
+            session.telemetry_history = history
+            session.history = history
+            session.selected_port = port
+
+            def telemetry_loop():
+                while not self.shutdown_event.is_set():
+                    try:
+                        self.datasource.select_device(uid)
+                        snap = self.datasource.snapshot()
+                        data = (snap.get("sensor_data")
+                                or snap.get("all_telemetry", {}).get(uid)
+                                or {})
+                        if data:
+                            telemetry_step(
+                                ctx, device_info=device_info,
+                                sensor_struct=data)
+                    except Exception as e:
+                        logger.warning(f"Telemetry error on {port}: {e}")
+                    time.sleep(0.1)
+
+            threading.Thread(target=telemetry_loop, daemon=True).start()
+
+    # ----------------- SESSION MANAGEMENT ----------------- #
+
+    def detect_and_start_sessions(self):
+        logger.info("Looking for WigiDash devices ...")
+        time.sleep(1)
+
+        usb_devices = scan_wigidash(self.vendor_id, self.product_id)
+        if not usb_devices:
+            logger.warning("No WigiDash devices detected.")
+            return
+
+        used_serials = {
+            s.usb_device.serial for s in self.sessions if s.usb_device.serial}
+
+        new_devices = []
+        for usb in usb_devices:
+            if not usb.serial:
+                try:
+                    usb.serial = usb.util.get_string(
+                        usb.dev, usb.dev.iSerialNumber)
+                except Exception as e:
+                    logger.warning(f"Failed to read USB serial: {e}")
+                    continue
+            if usb.serial in used_serials:
+                continue
+            new_devices.append(usb)
+
+        if not new_devices:
+            logger.info("No new WigiDash devices.")
+            return
+
+        logger.info("Scanning Benchlab devices ...")
+        time.sleep(1)
+        self.get_available_benchlabs()  # only builds cache, closes ports
+
+        logger.info("Starting Wigi sessions ...")
+        for usb in new_devices:
+            session = BenchlabWigiSession(
+                usb_device=usb,
+                telemetry_history=None,
+                manager=self,
+            )
+            try:
+                threading.Thread(target=session.run, daemon=True).start()
+                self.sessions.append(session)
+                used_serials.add(usb.serial)
+                logger.info(f"Started Wigidash session for {usb.serial}")
+            except Exception as e:
+                logger.error(f"Failed to start session for {usb.serial}: {e}")
+
+    def shutdown_manager(self):
+        """Close USB and COM ports safely after all sessions done."""
+        logger.info("Manager Shutdown Initiated.")
+
+        # Close COM ports
+        for port in self.benchlab_devices:
+            self.release_port(port)
+
+        self.sessions.clear()
+        logger.info("Manager shutdown completed")
+
+    def graceful_shutdown(self):
+        """Graceful shutdown: synchronize splash, then cleanup."""
+        if getattr(self, "shutting_down", False):
+            logger.info("Graceful shutdown already in progress.")
+            return
+
+        logger.info("Initiating graceful shutdown...")
+        self.shutting_down = True
+        self.shutdown_event.set()
+
+        # Disconnect DataSource if we created one
+        if self.datasource is not None:
+            try:
+                self.datasource.disconnect()
+                logger.info("DataSource disconnected")
+            except Exception as e:
+                logger.warning(f"Error disconnecting DataSource: {e}")
+
+        active_sessions = [s for s in self.sessions if s.app_running]
+
+        # Only create a barrier if we want splash synchronization
+        if active_sessions:
+            self.shutdown_barrier = threading.Barrier(len(active_sessions) + 1)
+
+        # Trigger shutdown in all sessions
+        for session in active_sessions:
+            threading.Thread(
+                target=session.shutdown_session,
+                daemon=True).start()
+
+        # If barrier exists, wait for synchronized splash
+        if self.shutdown_barrier:
+            try:
+                self.shutdown_barrier.wait()
+            except threading.BrokenBarrierError:
+                logger.warning("Shutdown barrier broken in manager.")
+            self.shutdown_barrier = None
+
+        # Wait for all sessions to signal cleanup_done
+        for session in active_sessions:
+            # blocks until each session finishes cleanup
+            session.cleanup_done.wait()
+
+        # Clean up manager
+        self.shutdown_manager()
+        logger.info("Graceful shutdown completed.")
+
+
+def main(args=None):
+    """Entry point. Accepts standard benchlab args namespace."""
+    if args is None:
+        args = types.SimpleNamespace(
+            source=os.environ.get(
+                "BENCHLAB_DATA_SOURCE",
+                "direct"),
+            api_url=os.environ.get(
+                "BENCHLAB_API_URL",
+                "http://127.0.0.1:8000"),
+            mqtt_broker=os.environ.get(
+                "MQTT_BROKER",
+                "localhost"),
+            mqtt_port=int(
+                os.environ.get(
+                    "MQTT_PORT",
+                    "1883")),
+            interval=1.0,
+        )
+
+    source = args.source
+    ds_kwargs = {}
+    if source in ("fastapi", "fastapi_custom"):
+        ds_kwargs["base_url"] = args.api_url
+        logger.info(f"Using FastAPI datasource: {args.api_url}")
+    elif source == "mqtt":
+        ds_kwargs["broker"] = args.mqtt_broker
+        ds_kwargs["port"] = args.mqtt_port
+        logger.info(
+            f"Using MQTT datasource: {
+                args.mqtt_broker}:{
+                args.mqtt_port}")
+    else:
+        logger.info("Using direct datasource")
+
+    datasource = DataSourceManager(source_type=source, **ds_kwargs)
+    if not datasource.connect():
+        logger.error(f"Failed to connect to {source} datasource")
+        return
+
+    manager = WigidashManager(datasource=datasource)
+    manager.detect_and_start_sessions()
+
+    try:
+        while True:
+            if manager.shutting_down:
+                break
+            if manager.sessions and all(
+                    s.cleanup_done.is_set() for s in manager.sessions):
+                break
+            time.sleep(1)
+    except KeyboardInterrupt:
+        logger.info("Initiating graceful shutdown")
+        manager.graceful_shutdown()
+    finally:
+        datasource.disconnect()
+
+
+if __name__ == "__main__":
+    main()

benchlab/wigidash/wigidash_session.py

@@ -0,0 +1,355 @@
+# wigidash_session.py
+
+from PIL import Image
+
+import os
+import time
+import threading
+
+from benchlab.wigidash.benchlab_fleet import BenchlabFleetSelect
+from benchlab.wigidash.benchlab_overview import BenchlabOverview
+from benchlab.wigidash.benchlab_graph import BenchlabGraph
+from benchlab.wigidash.benchlab_ui import UITheme
+from benchlab.wigidash.benchlab_utils import (
+    display_image, get_logger, KeepAliveManager
+)
+
+from benchlab.wigidash.wigidash_device import WigidashDevice
+from benchlab.wigidash.wigidash_widget import WidgetConfig
+
+
+logger = get_logger("WigidashSession")
+
+
+class BenchlabWigiSession:
+    """
+    UI session for a single Wigidash device.
+    """
+
+    SCREEN_WIDTH = UITheme.SCREEN_WIDTH
+    SCREEN_HEIGHT = UITheme.SCREEN_HEIGHT
+    SPLASH_TIME = 3.0
+    SCREEN_KEEPALIVE_INTERVAL = 5.0
+
+    def __init__(self, usb_device, telemetry_history=None, manager=None):
+        self.usb_device = usb_device
+        self.telemetry_history = telemetry_history
+        self.history = telemetry_history
+        self.telemetry_context = None
+        self.manager = manager
+        self.shutdown_event = threading.Event()
+
+        self.device_info = None
+        self.ser = None
+        self.uid = None
+        self.selected_port = None
+        self.sensor_data = {}
+
+        self.wigidash = None
+        self.keepalive_manager = None
+        self.app_running = True
+        self.cleanup_done = threading.Event()
+
+        self.fleet_page = None
+        self.overview_page = None
+        self.graph_page = None
+        self.graph_metrics = []
+
+        self.sensor_units = {
+            # --- High-level power ---
+            "SYS_Power": "W",
+            "CPU_Power": "W",
+            "GPU_Power": "W",
+            "MB_Power": "W",
+
+            # --- EPS Rails ---
+            "EPS1_Voltage": "V",
+            "EPS1_Current": "A",
+            "EPS1_Power": "W",
+            "EPS2_Voltage": "V",
+            "EPS2_Current": "A",
+            "EPS2_Power": "W",
+
+            # --- ATX Rails ---
+            "12V_Voltage": "V",
+            "12V_Current": "A",
+            "12V_Power": "W",
+            "5V_Voltage": "V",
+            "5V_Current": "A",
+            "5V_Power": "W",
+            "5VSB_Voltage": "V",
+            "5VSB_Current": "A",
+            "5VSB_Power": "W",
+            "3.3V_Voltage": "V",
+            "3.3V_Current": "A",
+            "3.3V_Power": "W",
+
+            # --- PCIe Rails ---
+            "PCIE8_1_Voltage": "V",
+            "PCIE8_1_Current": "A",
+            "PCIE8_1_Power": "W",
+            "PCIE8_2_Voltage": "V",
+            "PCIE8_2_Current": "A",
+            "PCIE8_2_Power": "W",
+            "PCIE8_3_Voltage": "V",
+            "PCIE8_3_Current": "A",
+            "PCIE8_3_Power": "W",
+            "HPWR1_Voltage": "V",
+            "HPWR1_Current": "A",
+            "HPWR1_Power": "W",
+            "HPWR2_Voltage": "V",
+            "HPWR2_Current": "A",
+            "HPWR2_Power": "W",
+
+            # --- VIN ---
+            # VIN_0 .. VIN_(SENSOR_VIN_NUM-1)
+            # safe default—adjust if SENSOR_VIN_NUM changes
+            **{f"VIN_{i}": "V" for i in range(16)},
+
+            # --- Other voltages ---
+            "Vdd": "V",
+            "Vref": "V",
+
+            # --- Temps ---
+            "Chip_Temp": "°C",
+            "Ambient_Temp": "°C",
+            "Humidity": "%",
+
+            # Temp sensors
+            **{f"Temp_Sensor_{i + 1}": "°C" for i in range(16)},
+
+            # --- Fans ---
+            **{f"Fan{i + 1}_Duty": "%" for i in range(16)},
+            **{f"Fan{i + 1}_RPM": "RPM" for i in range(16)},
+            # logical flag, unitless
+            **{f"Fan{i + 1}_Status": "" for i in range(16)},
+
+            "FanExtDuty": "%"
+        }
+
+    # ----------------- INIT ----------------- #
+    def connect_wigidash(self):
+        """Initialize the Wigidash device and UI."""
+        logger.info(
+            f"Initializing Wigidash device {
+                self.usb_device.serial} ...")
+        try:
+            self.usb_device.connect()
+            self.wigidash = WigidashDevice(self.usb_device)
+            self.wigidash.init_device()
+            self.wigidash.clear_page(0)
+            self.wigidash.change_page(0)
+
+            widget = WidgetConfig.create_fullscreen()
+            self.wigidash.add_widget(widget)
+
+            self.keepalive_manager = KeepAliveManager(
+                self.wigidash, interval=self.SCREEN_KEEPALIVE_INTERVAL)
+            self.keepalive_manager.start()
+
+            logger.info(
+                f"Wigidash device {
+                    self.usb_device.serial} initialized successfully.")
+            return True
+        except Exception as e:
+            logger.error(
+                f"Failed to initialize Wigidash {
+                    self.usb_device.serial}: {e}")
+            try:
+                self.usb_device.disconnect()
+            except Exception:
+                pass
+            return False
+
+    # ----------------- SPLASH ----------------- #
+
+    def show_splash(self):
+        """Display splash screen with logo."""
+        img = Image.new(
+            'RGB', (self.SCREEN_WIDTH, self.SCREEN_HEIGHT), color=(
+                41, 39, 38))
+        try:
+            base_dir = os.path.dirname(os.path.abspath(__file__))
+            assets_path = os.path.join(base_dir, "assets", "benchlab.png")
+            logo = Image.open(assets_path).convert('RGBA')
+            logo.thumbnail((500, 500), Image.Resampling.LANCZOS)
+            x = (self.SCREEN_WIDTH - logo.width) // 2
+            y = (self.SCREEN_HEIGHT - logo.height) // 2
+            img.paste(logo, (x, y), logo)
+        except Exception as e:
+            logger.warning(f"Failed to load splash logo: {e}")
+
+        display_image(self.wigidash, img)
+        logger.info(f"Splash screen displayed on {self.usb_device.serial}.")
+        time.sleep(self.SPLASH_TIME)
+
+    # ----------------- HELPER ----------------- #
+
+    def create_overview_page(self):
+        ov = BenchlabOverview(
+            self.wigidash,
+            wigi=self,
+            telemetry_history=self.telemetry_history,
+            telemetry_context=self.telemetry_context,
+            manager=self.manager
+        )
+        ov.start()
+        return ov
+
+    # ----------------- MAIN LOOP ----------------- #
+
+    def run(self):
+        """Run UI loop for Wigidash pages."""
+        logger.info(
+            f"Session {
+                self.usb_device.serial} started with manager={
+                self.manager}")
+        if not self.connect_wigidash():
+            return
+
+        self.show_splash()
+        self.next_page = "fleet"
+
+        while self.app_running and not self.shutdown_event.is_set():
+            try:
+                touch, _, _ = self.wigidash.get_click_info()
+            except Exception:
+                touch = None
+
+            if not self.app_running or self.shutdown_event.is_set():
+                break
+
+            # --- Fleet Page ---
+            if self.next_page == "fleet":
+                if self.fleet_page is None:
+                    self.fleet_page = BenchlabFleetSelect(
+                        self.wigidash,
+                        fleet_manager=self.manager,
+                        wigi=self
+                    )
+                    self.fleet_page.start()
+                    logger.info(
+                        f"Fleet page started for {
+                            self.usb_device.serial}.")
+
+                self.fleet_page.check_touch(touch)
+                self.fleet_page.render_and_display()
+
+                if (not self.fleet_page.running
+                        and self.fleet_page.selected_port):
+                    # Copy the selected port from fleet page into the session
+                    self.selected_port = self.fleet_page.selected_port
+
+                    # Fetch the telemetry context for that port
+                    if self.manager:
+                        self.telemetry_context = (
+                            self.manager.telemetry_contexts.get(
+                                self.selected_port))
+
+                    # Now create the overview page with a valid context
+                    self.overview_page = BenchlabOverview(
+                        self.wigidash,
+                        wigi=self,
+                        telemetry_history=self.telemetry_history,
+                        telemetry_context=self.telemetry_context,
+                        manager=self.manager or self
+                    )
+                    self.overview_page.start()
+                    self.next_page = "overview"
+                    self.fleet_page = None
+                    logger.info(
+                        f"Fleet selection done on {
+                            self.usb_device.serial}, opening Overview.")
+
+            # --- Overview Page ---
+            elif self.next_page == "overview":
+                if self.overview_page:
+                    self.overview_page.check_touch(touch)
+                    self.overview_page.render_and_display()
+
+                    if not self.overview_page.running:
+                        requested_metrics = getattr(
+                            self.overview_page, "requested_graph_metrics", [])
+                        self.overview_page = None
+
+                        if requested_metrics:
+                            self.graph_metrics = requested_metrics
+                            self.next_page = "graph"
+                            logger.info(
+                                f"Overview requested graph metrics: {
+                                    self.graph_metrics}")
+                        else:
+                            self.next_page = "fleet"
+
+            # --- Graph Page ---
+            elif self.next_page == "graph":
+                if self.graph_page is None:
+                    self.graph_page = BenchlabGraph(
+                        self.wigidash,
+                        self,
+                        metrics=self.graph_metrics,
+                        telemetry_history=self.telemetry_history,
+                        telemetry_context=self.telemetry_context,
+                        manager=self.manager or self
+                    )
+                    self.graph_page.start()
+                    logger.info(
+                        f"Graph page started for {
+                            self.usb_device.serial} with metrics {
+                            self.graph_metrics}.")
+
+                self.graph_page.check_touch(touch)
+                self.graph_page.render_and_display()
+
+                if not self.graph_page.running:
+                    self.graph_page = None
+                    self.next_page = "overview"
+                    logger.info(
+                        f"Graph page exited for {
+                            self.usb_device.serial}, returning to Overview.")
+
+            time.sleep(0.05)
+
+    # ----------------- STOP / CLEANUP ----------------- #
+
+    def cleanup(self):
+        """Stop all pages and keepalive manager."""
+        if self.keepalive_manager:
+            self.keepalive_manager.stop()
+        if self.wigidash:
+            try:
+                self.wigidash.clear_page(0)
+            except Exception:
+                pass
+        if self.usb_device:
+            try:
+                self.usb_device.disconnect()
+            except Exception as e:
+                logger.warning(f"Error disconnecting USB device: {e}")
+        logger.info(f"Wigidash session {self.usb_device.serial} cleaned up.")
+
+    def shutdown_session(self):
+        """Show synchronized shutdown splash, then stop pages/keepalive."""
+        logger.info(f"Stopping Wigidash session {self.usb_device.serial}")
+        self.app_running = False
+        self.shutdown_event.set()
+
+        # Wait at barrier if it exists
+        barrier = getattr(self.manager, "shutdown_barrier", None)
+        if barrier:
+            try:
+                barrier.wait()
+            except threading.BrokenBarrierError:
+                logger.warning(
+                    f"Shutdown barrier broken for {
+                        self.usb_device.serial}")
+
+        # Show splash
+        if self.wigidash:
+            try:
+                self.show_splash()
+            except Exception as e:
+                logger.warning(f"Failed to show shutdown splash: {e}")
+
+        self.cleanup()
+        self.cleanup_done.set()