From f67c794896aefda012ed4c42120fb0fc6c7526db Mon Sep 17 00:00:00 2001 From: Rick Wierenga Date: Wed, 12 Aug 2026 13:59:13 -0700 Subject: [PATCH 1/3] fix(io): run blocking device calls on the io executor, not the event loop FTDI's write, read and readline called into libftdi directly from async methods, blocking the event loop for the duration of the transfer. Move them onto the single worker that the io object already owns. Device open and close had the same problem in ftdi, hid and usb, and the serial port scan in serial: enumeration, configuration and the usb read buffer drain all ran on the loop. Each is now a _setup_sync that runs as one unit on the executor, so the handle is opened by the thread that later uses it and no other coroutine observes a half-configured device. Executors are created before the open and torn down when it fails, which also fixes a thread leak in serial's setup, where the "no machines found" and "multiple devices detected" paths returned without shutting theirs down. Shutdown is now non-blocking: the worker is idle by that point, so waiting on it could only stall the loop behind an unrelated queued read. Co-Authored-By: Claude Opus 5 (1M context) --- pylabrobot/io/ftdi.py | 65 ++++++++++++++++++++++++++--------------- pylabrobot/io/hid.py | 45 ++++++++++++++++++---------- pylabrobot/io/serial.py | 43 ++++++++++++++++----------- pylabrobot/io/usb.py | 58 +++++++++++++++++++++++------------- 4 files changed, 136 insertions(+), 75 deletions(-) diff --git a/pylabrobot/io/ftdi.py b/pylabrobot/io/ftdi.py index 17df02e7603..b504ded5f8a 100644 --- a/pylabrobot/io/ftdi.py +++ b/pylabrobot/io/ftdi.py @@ -178,32 +178,48 @@ def _resolve_device_serial(self) -> str: device_serial_number = cast(str, usb.util.get_string(device, device.iSerialNumber)) return device_serial_number - async def setup(self): - """Initialize the FTDI device connection with device resolution.""" + def _setup_sync(self) -> None: + """Resolve and open the device. Runs on the executor that owns all device calls.""" if self._dev is not None and not self._dev.closed: self._dev.close() + + # Resolve which device to connect to + self._device_id = self._resolve_device_serial() + + # Create and open device + self._dev = Device( + lazy_open=True, + device_id=self.device_id, + pid=self._pid, + vid=self._vid, + interface_select=self._interface_select, + ) + self._dev.open() + + async def setup(self): + """Initialize the FTDI device connection with device resolution.""" + if self._executor is None: + self._executor = ThreadPoolExecutor(max_workers=1) + loop = asyncio.get_running_loop() try: - # Resolve which device to connect to - self._device_id = self._resolve_device_serial() - - # Create and open device - self._dev = Device( - lazy_open=True, - device_id=self.device_id, - pid=self._pid, - vid=self._vid, - interface_select=self._interface_select, - ) - self._dev.open() - logger.info(f"Successfully opened FTDI device: {self.device_id}") + await loop.run_in_executor(self._executor, self._setup_sync) except FtdiError as e: + self._shutdown_executor() raise RuntimeError( f"Failed to open FTDI device for '{self.human_readable_device_name}': {e}. " "Is the device connected? Is it in use by another process? " "Try restarting the kernel." ) from e + except BaseException: + self._shutdown_executor() + raise + logger.info(f"Successfully opened FTDI device: {self.device_id}") - self._executor = ThreadPoolExecutor(max_workers=1) + def _shutdown_executor(self) -> None: + if self._executor is not None: + # the worker is idle here, so this does not block the event loop + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None @property def device_id(self) -> str: @@ -297,20 +313,22 @@ async def request_serial(self) -> str: return self.device_id async def stop(self): + loop = asyncio.get_running_loop() if self._dev is not None: - self.dev.close() - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None + await loop.run_in_executor(self._executor, self.dev.close) + self._dev = None + self._shutdown_executor() async def write(self, data: bytes) -> int: """Write data to the device. Returns the number of bytes written.""" logger.log(LOG_LEVEL_IO, "[%s] write %s", self._device_id, data) capturer.record(FTDICommand(device_id=self.device_id, action="write", data=data.hex())) - return cast(int, self.dev.write(data)) + loop = asyncio.get_running_loop() + return cast(int, await loop.run_in_executor(self._executor, self.dev.write, data)) async def read(self, num_bytes: int = 1) -> bytes: - data = self.dev.read(num_bytes) + loop = asyncio.get_running_loop() + data = await loop.run_in_executor(self._executor, self.dev.read, num_bytes) if len(data) != 0: logger.log(LOG_LEVEL_IO, "[%s] read %s", self._device_id, data) capturer.record( @@ -323,7 +341,8 @@ async def read(self, num_bytes: int = 1) -> bytes: return cast(bytes, data) async def readline(self) -> bytes: # type: ignore # very dumb it's reading from pyserial - data = self.dev.readline() + loop = asyncio.get_running_loop() + data = await loop.run_in_executor(self._executor, self.dev.readline) if len(data) != 0: logger.log(LOG_LEVEL_IO, "[%s] readline %s", self._device_id, data) capturer.record(FTDICommand(device_id=self.device_id, action="readline", data=data.hex())) diff --git a/pylabrobot/io/hid.py b/pylabrobot/io/hid.py index 1b624da6962..fcb2821bd18 100644 --- a/pylabrobot/io/hid.py +++ b/pylabrobot/io/hid.py @@ -43,17 +43,8 @@ def __init__( if get_capture_or_validation_active(): raise RuntimeError("Cannot create a new HID object while capture or validation is active") - async def setup(self): - """ - Sets up the HID device by enumerating connected devices, matching the specified - VID, PID, and optional serial number, and opening a connection to the device. - """ - if not USE_HID: - raise RuntimeError( - "hid is not installed. Install with: pip install pylabrobot[hid]. " - f"Import error: {_HID_IMPORT_ERROR}" - ) - + def _setup_sync(self) -> None: + """Enumerate, match and open the device. Runs on the executor that owns all device calls.""" # --- 1. Enumerate all HID devices --- all_devices = hid.enumerate() candidates = [ @@ -102,19 +93,43 @@ async def setup(self): self.device = hid.Device( path=chosen["path"] # safer than vid/pid/serial triple ) + + async def setup(self): + """ + Sets up the HID device by enumerating connected devices, matching the specified + VID, PID, and optional serial number, and opening a connection to the device. + """ + if not USE_HID: + raise RuntimeError( + "hid is not installed. Install with: pip install pylabrobot[hid]. " + f"Import error: {_HID_IMPORT_ERROR}" + ) + self._executor = ThreadPoolExecutor(max_workers=1) + loop = asyncio.get_running_loop() + try: + await loop.run_in_executor(self._executor, self._setup_sync) + except BaseException: + self._shutdown_executor() + raise logger.log(LOG_LEVEL_IO, "Opened HID device %s", self._unique_id) capturer.record(HIDCommand(device_id=self._unique_id, action="open", data="")) + def _shutdown_executor(self) -> None: + if self._executor is not None: + # the worker is idle here, so this does not block the event loop + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None + async def stop(self): if self.device is not None: - self.device.close() + loop = asyncio.get_running_loop() + await loop.run_in_executor(self._executor, self.device.close) + self.device = None logger.log(LOG_LEVEL_IO, "Closing HID device %s", self._unique_id) capturer.record(HIDCommand(device_id=self._unique_id, action="close", data="")) - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None + self._shutdown_executor() async def write(self, data: bytes, report_id: bytes = b"\x00"): r"""Writes data to the HID device. diff --git a/pylabrobot/io/serial.py b/pylabrobot/io/serial.py index 26ddf801d96..6095008ffe1 100644 --- a/pylabrobot/io/serial.py +++ b/pylabrobot/io/serial.py @@ -147,6 +147,26 @@ async def setup(self): loop = asyncio.get_running_loop() self._executor = ThreadPoolExecutor(max_workers=1) + try: + self._port = await loop.run_in_executor(self._executor, self._setup_sync) + except BaseException: + self._shutdown_executor() + raise + + assert self._ser is not None + + def _shutdown_executor(self) -> None: + if self._executor is not None: + # the worker is idle here, so this does not block the event loop + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None + + def _setup_sync(self) -> str: + """Resolve the port and open the connection. Runs on the executor that owns all device calls. + + Returns the resolved port. + """ + # 1. VID:PID specified - port maybe if self._vid is not None and self._pid is not None: matching_ports = [ @@ -183,8 +203,8 @@ async def setup(self): "Please specify the correct port address explicitly (e.g. /dev/ttyUSB0 or COM3)." ) - def _open_serial() -> serial.Serial: - return serial.Serial( + try: + self._ser = serial.Serial( port=candidate_port, baudrate=self.baudrate, bytesize=self.bytesize, @@ -196,22 +216,13 @@ def _open_serial() -> serial.Serial: dsrdtr=self.dsrdtr, xonxoff=self.xonxoff, ) - - try: - self._ser = await loop.run_in_executor(self._executor, _open_serial) - - except serial.SerialException as e: + except serial.SerialException: logger.error( f"Could not connect to device '{self.human_readable_device_name}', is it in use by a different notebook/process?" ) - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None - raise e + raise - assert self._ser is not None - - self._port = candidate_port + return candidate_port async def stop(self): """Close the serial device.""" @@ -223,9 +234,7 @@ async def stop(self): raise RuntimeError(f"Call setup() first for device '{self.human_readable_device_name}'.") await loop.run_in_executor(self._executor, self._ser.close) - if self._executor is not None: - self._executor.shutdown(wait=True) - self._executor = None + self._shutdown_executor() async def write(self, data: bytes): """Write data to the serial device.""" diff --git a/pylabrobot/io/usb.py b/pylabrobot/io/usb.py index 8e4469f1577..fb6a4eab42d 100644 --- a/pylabrobot/io/usb.py +++ b/pylabrobot/io/usb.py @@ -383,21 +383,8 @@ def ctrl_transfer( return bytearray(res) - async def setup(self, empty_buffer=True): - """Initialize the USB connection to the machine.""" - - if self.dev is not None: - # previous setup did not properly finish, - # or we are re-initializing the device. - logger.warning("USB device already connected. Closing previous connection.") - await self.stop() - - if not USE_USB: - raise RuntimeError( - "pyusb/libusb is not installed. Install with: pip install pylabrobot[usb]. " - f"Import error: {_USB_IMPORT_ERROR}. " - "https://docs.pylabrobot.org/installation.html" - ) + def _setup_sync(self, empty_buffer: bool) -> None: + """Open the device, resolve the endpoints and drain stale packets. Runs off the event loop.""" logger.info("Finding USB device...") @@ -457,22 +444,53 @@ async def setup(self, empty_buffer=True): while self._read_packet() is not None: pass + async def setup(self, empty_buffer=True): + """Initialize the USB connection to the machine.""" + + if self.dev is not None: + # previous setup did not properly finish, + # or we are re-initializing the device. + logger.warning("USB device already connected. Closing previous connection.") + await self.stop() + + if not USE_USB: + raise RuntimeError( + "pyusb/libusb is not installed. Install with: pip install pylabrobot[usb]. " + f"Import error: {_USB_IMPORT_ERROR}. " + "https://docs.pylabrobot.org/installation.html" + ) + self._read_executor = ThreadPoolExecutor(max_workers=self.max_workers) self._write_executor = ThreadPoolExecutor(max_workers=self.max_workers) + loop = asyncio.get_running_loop() + try: + await loop.run_in_executor(self._read_executor, self._setup_sync, empty_buffer) + except BaseException: + self._shutdown_executors() + self.dev = None + raise + + def _shutdown_executors(self) -> None: + for executor in (self._read_executor, self._write_executor): + if executor is not None: + # the workers are idle here, so this does not block the event loop + executor.shutdown(wait=False, cancel_futures=True) + self._read_executor = self._write_executor = None + async def stop(self): """Close the USB connection to the machine. Safe to call multiple times.""" if self.dev is None: + self._shutdown_executors() return logger.warning("Closing connection to USB device.") - usb.util.dispose_resources(self.dev) + loop = asyncio.get_running_loop() + dev = self.dev + await loop.run_in_executor(self._read_executor, lambda: usb.util.dispose_resources(dev)) self.dev = None - for executor in (self._read_executor, self._write_executor): - if executor is not None: - executor.shutdown(wait=True) - self._read_executor = self._write_executor = None + self._shutdown_executors() def serialize(self) -> dict: """Serialize the backend to a dictionary.""" From 6ee7ad2d3c3db3a25eb9123325582b9b65253938 Mon Sep 17 00:00:00 2001 From: Rick Wierenga Date: Thu, 13 Aug 2026 12:46:28 -0700 Subject: [PATCH 2/3] fix(io): clean up cancelled device setup --- pylabrobot/io/executor_setup_tests.py | 194 ++++++++++++++++++++++++++ pylabrobot/io/ftdi.py | 37 +++-- pylabrobot/io/hid.py | 25 +++- pylabrobot/io/io.py | 25 ++++ pylabrobot/io/serial.py | 25 +++- pylabrobot/io/usb.py | 31 +++- 6 files changed, 315 insertions(+), 22 deletions(-) create mode 100644 pylabrobot/io/executor_setup_tests.py diff --git a/pylabrobot/io/executor_setup_tests.py b/pylabrobot/io/executor_setup_tests.py new file mode 100644 index 00000000000..8947f3343d1 --- /dev/null +++ b/pylabrobot/io/executor_setup_tests.py @@ -0,0 +1,194 @@ +import asyncio +import threading +import unittest +from types import SimpleNamespace +from typing import Any, Callable, Coroutine, Optional +from unittest import mock + +from pylabrobot.io import ftdi as ftdi_module +from pylabrobot.io import hid as hid_module +from pylabrobot.io import serial as serial_module +from pylabrobot.io import usb as usb_module + + +class _ClosableDevice: + def __init__(self) -> None: + self.setup_thread: Optional[int] = None + self.close_thread: Optional[int] = None + + def close(self) -> None: + self.close_thread = threading.get_ident() + + +class _DisposableDevice: + def __init__(self) -> None: + self.setup_thread: Optional[int] = None + self.dispose_thread: Optional[int] = None + + +class ExecutorSetupTests(unittest.IsolatedAsyncioTestCase): + async def _wait_for_task(self, task: "asyncio.Task[None]") -> None: + while not task.done(): + await asyncio.wait({task}, timeout=0.01) + await task + + async def _cancel_while_worker_is_running( + self, + setup: Callable[[], Coroutine[Any, Any, None]], + started: threading.Event, + release: threading.Event, + ) -> None: + task: "asyncio.Task[None]" = asyncio.create_task(setup()) + while not started.is_set(): + await asyncio.sleep(0) + + try: + task.cancel() + await asyncio.sleep(0) + self.assertFalse(task.done(), "setup returned before its worker finished") + finally: + release.set() + + with self.assertRaises(asyncio.CancelledError): + await self._wait_for_task(task) + + async def test_cancelled_ftdi_setup_closes_device_before_executor_shutdown(self) -> None: + io = ftdi_module.FTDI.__new__(ftdi_module.FTDI) + io.human_readable_device_name = "mock FTDI" + io._device_id = "mock" + io._dev = None + io._executor = None + device = _ClosableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync() -> None: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io._dev = device # type: ignore[assignment] + + io._setup_sync = setup_sync # type: ignore[method-assign] + ftdi_error = type("MockFtdiError", (Exception,), {}) + with mock.patch.object(ftdi_module, "FtdiError", ftdi_error, create=True): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.close_thread, device.setup_thread) + self.assertIsNone(io._dev) + self.assertIsNone(io._executor) + + async def test_cancelled_hid_setup_closes_device_before_executor_shutdown(self) -> None: + io = hid_module.HID.__new__(hid_module.HID) + io.human_readable_device_name = "mock HID" + io._unique_id = "mock" + io.device = None + io._executor = None + device = _ClosableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync() -> None: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io.device = device # type: ignore[assignment] + + io._setup_sync = setup_sync # type: ignore[method-assign] + with mock.patch.object(hid_module, "USE_HID", True): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.close_thread, device.setup_thread) + self.assertIsNone(io.device) + self.assertIsNone(io._executor) + + async def test_cancelled_serial_setup_closes_port_before_executor_shutdown(self) -> None: + io = serial_module.Serial.__new__(serial_module.Serial) + io.human_readable_device_name = "mock serial" + io._ser = None + io._executor = None + device = _ClosableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync() -> str: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io._ser = device # type: ignore[assignment] + return "/dev/mock" + + io._setup_sync = setup_sync # type: ignore[method-assign] + with mock.patch.object(serial_module, "HAS_SERIAL", True): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.close_thread, device.setup_thread) + self.assertIsNone(io._ser) + self.assertIsNone(io._executor) + + async def test_cancelled_usb_setup_disposes_device_before_executor_shutdown(self) -> None: + io = usb_module.USB( + id_vendor=1, + id_product=2, + human_readable_device_name="mock USB", + packet_read_timeout=1, + read_timeout=2, + ) + device = _DisposableDevice() + started = threading.Event() + release = threading.Event() + + def setup_sync(empty_buffer: bool) -> None: + device.setup_thread = threading.get_ident() + started.set() + release.wait() + io.dev = device # type: ignore[assignment] + + def dispose_resources(dev: object) -> None: + self.assertIs(dev, device) + device.dispose_thread = threading.get_ident() + + io._setup_sync = setup_sync # type: ignore[method-assign] + fake_usb = SimpleNamespace(util=SimpleNamespace(dispose_resources=dispose_resources)) + with ( + mock.patch.object(usb_module, "USE_USB", True), + mock.patch.object(usb_module, "usb", fake_usb, create=True), + ): + await self._cancel_while_worker_is_running(io.setup, started, release) + + self.assertEqual(device.dispose_thread, device.setup_thread) + self.assertIsNone(io.dev) + self.assertIsNone(io._read_executor) + self.assertIsNone(io._write_executor) + + async def test_usb_setup_error_disposes_acquired_device_on_worker(self) -> None: + io = usb_module.USB( + id_vendor=1, + id_product=2, + human_readable_device_name="mock USB", + packet_read_timeout=1, + read_timeout=2, + ) + device = _DisposableDevice() + + def setup_sync(empty_buffer: bool) -> None: + device.setup_thread = threading.get_ident() + io.dev = device # type: ignore[assignment] + raise RuntimeError("configuration failed") + + def dispose_resources(dev: object) -> None: + self.assertIs(dev, device) + device.dispose_thread = threading.get_ident() + + io._setup_sync = setup_sync # type: ignore[method-assign] + fake_usb = SimpleNamespace(util=SimpleNamespace(dispose_resources=dispose_resources)) + with ( + mock.patch.object(usb_module, "USE_USB", True), + mock.patch.object(usb_module, "usb", fake_usb, create=True), + self.assertRaisesRegex(RuntimeError, "configuration failed"), + ): + await self._wait_for_task(asyncio.create_task(io.setup())) + + self.assertEqual(device.dispose_thread, device.setup_thread) + self.assertIsNone(io.dev) + self.assertIsNone(io._read_executor) + self.assertIsNone(io._write_executor) diff --git a/pylabrobot/io/ftdi.py b/pylabrobot/io/ftdi.py index b504ded5f8a..e56a5af8cb9 100644 --- a/pylabrobot/io/ftdi.py +++ b/pylabrobot/io/ftdi.py @@ -25,6 +25,7 @@ from pylabrobot.io.capture import CaptureReader, Command, capturer, get_capture_or_validation_active from pylabrobot.io.errors import ValidationError +from pylabrobot.io.io import _wait_for_executor_future from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences logger = logging.getLogger(__name__) @@ -201,18 +202,34 @@ async def setup(self): if self._executor is None: self._executor = ThreadPoolExecutor(max_workers=1) loop = asyncio.get_running_loop() + setup_future = loop.run_in_executor(self._executor, self._setup_sync) + setup_error: Optional[BaseException] = None try: - await loop.run_in_executor(self._executor, self._setup_sync) - except FtdiError as e: - self._shutdown_executor() - raise RuntimeError( - f"Failed to open FTDI device for '{self.human_readable_device_name}': {e}. " - "Is the device connected? Is it in use by another process? " - "Try restarting the kernel." - ) from e - except BaseException: + await asyncio.shield(setup_future) + except BaseException as exc: + setup_error = exc + if not setup_future.done(): + try: + await _wait_for_executor_future(setup_future) + except BaseException: + pass + + if setup_error is not None: + if self._dev is not None and self._executor is not None: + close_future = loop.run_in_executor(self._executor, self._dev.close) + try: + await _wait_for_executor_future(close_future) + except Exception: + logger.warning("Failed to close FTDI device after setup failure", exc_info=True) + self._dev = None self._shutdown_executor() - raise + if isinstance(setup_error, FtdiError): + raise RuntimeError( + f"Failed to open FTDI device for '{self.human_readable_device_name}': {setup_error}. " + "Is the device connected? Is it in use by another process? " + "Try restarting the kernel." + ) from setup_error + raise setup_error logger.info(f"Successfully opened FTDI device: {self.device_id}") def _shutdown_executor(self) -> None: diff --git a/pylabrobot/io/hid.py b/pylabrobot/io/hid.py index fcb2821bd18..5754ae453f9 100644 --- a/pylabrobot/io/hid.py +++ b/pylabrobot/io/hid.py @@ -5,7 +5,7 @@ from pylabrobot.io.capture import CaptureReader, Command, capturer, get_capture_or_validation_active from pylabrobot.io.errors import ValidationError -from pylabrobot.io.io import IOBase +from pylabrobot.io.io import IOBase, _wait_for_executor_future from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences try: @@ -107,11 +107,28 @@ async def setup(self): self._executor = ThreadPoolExecutor(max_workers=1) loop = asyncio.get_running_loop() + setup_future = loop.run_in_executor(self._executor, self._setup_sync) + setup_error: Optional[BaseException] = None try: - await loop.run_in_executor(self._executor, self._setup_sync) - except BaseException: + await asyncio.shield(setup_future) + except BaseException as exc: + setup_error = exc + if not setup_future.done(): + try: + await _wait_for_executor_future(setup_future) + except BaseException: + pass + + if setup_error is not None: + if self.device is not None and self._executor is not None: + close_future = loop.run_in_executor(self._executor, self.device.close) + try: + await _wait_for_executor_future(close_future) + except Exception: + logger.warning("Failed to close HID device after setup failure", exc_info=True) + self.device = None self._shutdown_executor() - raise + raise setup_error logger.log(LOG_LEVEL_IO, "Opened HID device %s", self._unique_id) capturer.record(HIDCommand(device_id=self._unique_id, action="open", data="")) diff --git a/pylabrobot/io/io.py b/pylabrobot/io/io.py index 599399a251f..ddd269e1b10 100644 --- a/pylabrobot/io/io.py +++ b/pylabrobot/io/io.py @@ -1,8 +1,33 @@ +import asyncio from abc import ABC, abstractmethod +from typing import TypeVar from pylabrobot.serializer import SerializableMixin +T = TypeVar("T") + + +async def _wait_for_executor_future(future: "asyncio.Future[T]") -> T: + """Wait for executor work to finish, ignoring cancellation of the awaiting task. + + Cancelling the asyncio wrapper returned by ``run_in_executor`` does not stop a callable that is + already running. This keeps teardown from racing work that still owns a device handle. + """ + + async def wait() -> T: + return await future + + waiter = asyncio.create_task(wait()) + while not waiter.done(): + try: + await asyncio.shield(waiter) + except asyncio.CancelledError: + continue + + return waiter.result() + + class IOBase(SerializableMixin, ABC): @abstractmethod async def write(self, data: bytes, *args, **kwargs): diff --git a/pylabrobot/io/serial.py b/pylabrobot/io/serial.py index 6095008ffe1..590a102416f 100644 --- a/pylabrobot/io/serial.py +++ b/pylabrobot/io/serial.py @@ -18,6 +18,7 @@ _SERIAL_IMPORT_ERROR = e from pylabrobot.io.capture import CaptureReader, Command, capturer, get_capture_or_validation_active +from pylabrobot.io.io import _wait_for_executor_future from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences logger = logging.getLogger(__name__) @@ -147,11 +148,29 @@ async def setup(self): loop = asyncio.get_running_loop() self._executor = ThreadPoolExecutor(max_workers=1) + setup_future = loop.run_in_executor(self._executor, self._setup_sync) + setup_error: Optional[BaseException] = None try: - self._port = await loop.run_in_executor(self._executor, self._setup_sync) - except BaseException: + resolved_port = await asyncio.shield(setup_future) + self._port = resolved_port + except BaseException as exc: + setup_error = exc + if not setup_future.done(): + try: + await _wait_for_executor_future(setup_future) + except BaseException: + pass + + if setup_error is not None: + if self._ser is not None and self._executor is not None: + close_future = loop.run_in_executor(self._executor, self._ser.close) + try: + await _wait_for_executor_future(close_future) + except Exception: + logger.warning("Failed to close serial port after setup failure", exc_info=True) + self._ser = None self._shutdown_executor() - raise + raise setup_error assert self._ser is not None diff --git a/pylabrobot/io/usb.py b/pylabrobot/io/usb.py index fb6a4eab42d..fea19116d06 100644 --- a/pylabrobot/io/usb.py +++ b/pylabrobot/io/usb.py @@ -7,7 +7,7 @@ from pylabrobot.io.capture import Command, capturer, get_capture_or_validation_active from pylabrobot.io.errors import ValidationError -from pylabrobot.io.io import IOBase +from pylabrobot.io.io import IOBase, _wait_for_executor_future from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences try: @@ -464,12 +464,33 @@ async def setup(self, empty_buffer=True): self._write_executor = ThreadPoolExecutor(max_workers=self.max_workers) loop = asyncio.get_running_loop() + setup_future = loop.run_in_executor(self._read_executor, self._setup_sync, empty_buffer) + setup_error: Optional[BaseException] = None try: - await loop.run_in_executor(self._read_executor, self._setup_sync, empty_buffer) - except BaseException: - self._shutdown_executors() + await asyncio.shield(setup_future) + except BaseException as exc: + setup_error = exc + if not setup_future.done(): + try: + await _wait_for_executor_future(setup_future) + except BaseException: + pass + + if setup_error is not None: + if self.dev is not None and self._read_executor is not None: + dev = self.dev + dispose_future = loop.run_in_executor( + self._read_executor, lambda: usb.util.dispose_resources(dev) + ) + try: + await _wait_for_executor_future(dispose_future) + except Exception: + logger.warning("Failed to dispose USB device after setup failure", exc_info=True) self.dev = None - raise + self.read_endpoint = None + self.write_endpoint = None + self._shutdown_executors() + raise setup_error def _shutdown_executors(self) -> None: for executor in (self._read_executor, self._write_executor): From fca07dcc9180f9e90874cfe429b565b28955a9dd Mon Sep 17 00:00:00 2001 From: Rick Wierenga Date: Thu, 13 Aug 2026 13:15:50 -0700 Subject: [PATCH 3/3] refactor(io): simplify cancelled setup cleanup --- pylabrobot/io/executor_setup_tests.py | 11 ++++---- pylabrobot/io/ftdi.py | 37 +++++++++++++++------------ pylabrobot/io/hid.py | 19 +++++--------- pylabrobot/io/io.py | 25 ------------------ pylabrobot/io/serial.py | 18 +++++-------- pylabrobot/io/usb.py | 25 +++++++----------- 6 files changed, 47 insertions(+), 88 deletions(-) diff --git a/pylabrobot/io/executor_setup_tests.py b/pylabrobot/io/executor_setup_tests.py index 8947f3343d1..ad6a2e02697 100644 --- a/pylabrobot/io/executor_setup_tests.py +++ b/pylabrobot/io/executor_setup_tests.py @@ -25,6 +25,10 @@ def __init__(self) -> None: self.setup_thread: Optional[int] = None self.dispose_thread: Optional[int] = None + def set_configuration(self) -> None: + self.setup_thread = threading.get_ident() + raise RuntimeError("configuration failed") + class ExecutorSetupTests(unittest.IsolatedAsyncioTestCase): async def _wait_for_task(self, task: "asyncio.Task[None]") -> None: @@ -170,16 +174,11 @@ async def test_usb_setup_error_disposes_acquired_device_on_worker(self) -> None: ) device = _DisposableDevice() - def setup_sync(empty_buffer: bool) -> None: - device.setup_thread = threading.get_ident() - io.dev = device # type: ignore[assignment] - raise RuntimeError("configuration failed") - def dispose_resources(dev: object) -> None: self.assertIs(dev, device) device.dispose_thread = threading.get_ident() - io._setup_sync = setup_sync # type: ignore[method-assign] + io.get_available_devices = mock.Mock(return_value=[device]) # type: ignore[method-assign] fake_usb = SimpleNamespace(util=SimpleNamespace(dispose_resources=dispose_resources)) with ( mock.patch.object(usb_module, "USE_USB", True), diff --git a/pylabrobot/io/ftdi.py b/pylabrobot/io/ftdi.py index e56a5af8cb9..53b9594c506 100644 --- a/pylabrobot/io/ftdi.py +++ b/pylabrobot/io/ftdi.py @@ -25,7 +25,6 @@ from pylabrobot.io.capture import CaptureReader, Command, capturer, get_capture_or_validation_active from pylabrobot.io.errors import ValidationError -from pylabrobot.io.io import _wait_for_executor_future from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences logger = logging.getLogger(__name__) @@ -183,19 +182,28 @@ def _setup_sync(self) -> None: """Resolve and open the device. Runs on the executor that owns all device calls.""" if self._dev is not None and not self._dev.closed: self._dev.close() + self._dev = None # Resolve which device to connect to self._device_id = self._resolve_device_serial() # Create and open device - self._dev = Device( + dev = Device( lazy_open=True, device_id=self.device_id, pid=self._pid, vid=self._vid, interface_select=self._interface_select, ) - self._dev.open() + try: + dev.open() + except BaseException: + try: + dev.close() + except Exception: + logger.warning("Failed to close FTDI device after setup failure", exc_info=True) + raise + self._dev = dev async def setup(self): """Initialize the FTDI device connection with device resolution.""" @@ -203,33 +211,28 @@ async def setup(self): self._executor = ThreadPoolExecutor(max_workers=1) loop = asyncio.get_running_loop() setup_future = loop.run_in_executor(self._executor, self._setup_sync) - setup_error: Optional[BaseException] = None try: await asyncio.shield(setup_future) except BaseException as exc: - setup_error = exc - if not setup_future.done(): + if isinstance(exc, asyncio.CancelledError): try: - await _wait_for_executor_future(setup_future) + await setup_future except BaseException: pass - - if setup_error is not None: - if self._dev is not None and self._executor is not None: - close_future = loop.run_in_executor(self._executor, self._dev.close) + if self._dev is not None: try: - await _wait_for_executor_future(close_future) + await loop.run_in_executor(self._executor, self._dev.close) except Exception: logger.warning("Failed to close FTDI device after setup failure", exc_info=True) - self._dev = None + self._dev = None self._shutdown_executor() - if isinstance(setup_error, FtdiError): + if isinstance(exc, FtdiError): raise RuntimeError( - f"Failed to open FTDI device for '{self.human_readable_device_name}': {setup_error}. " + f"Failed to open FTDI device for '{self.human_readable_device_name}': {exc}. " "Is the device connected? Is it in use by another process? " "Try restarting the kernel." - ) from setup_error - raise setup_error + ) from exc + raise logger.info(f"Successfully opened FTDI device: {self.device_id}") def _shutdown_executor(self) -> None: diff --git a/pylabrobot/io/hid.py b/pylabrobot/io/hid.py index 5754ae453f9..1a091522021 100644 --- a/pylabrobot/io/hid.py +++ b/pylabrobot/io/hid.py @@ -5,7 +5,7 @@ from pylabrobot.io.capture import CaptureReader, Command, capturer, get_capture_or_validation_active from pylabrobot.io.errors import ValidationError -from pylabrobot.io.io import IOBase, _wait_for_executor_future +from pylabrobot.io.io import IOBase from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences try: @@ -108,27 +108,22 @@ async def setup(self): self._executor = ThreadPoolExecutor(max_workers=1) loop = asyncio.get_running_loop() setup_future = loop.run_in_executor(self._executor, self._setup_sync) - setup_error: Optional[BaseException] = None try: await asyncio.shield(setup_future) except BaseException as exc: - setup_error = exc - if not setup_future.done(): + if isinstance(exc, asyncio.CancelledError): try: - await _wait_for_executor_future(setup_future) + await setup_future except BaseException: pass - - if setup_error is not None: - if self.device is not None and self._executor is not None: - close_future = loop.run_in_executor(self._executor, self.device.close) + if self.device is not None: try: - await _wait_for_executor_future(close_future) + await loop.run_in_executor(self._executor, self.device.close) except Exception: logger.warning("Failed to close HID device after setup failure", exc_info=True) - self.device = None + self.device = None self._shutdown_executor() - raise setup_error + raise logger.log(LOG_LEVEL_IO, "Opened HID device %s", self._unique_id) capturer.record(HIDCommand(device_id=self._unique_id, action="open", data="")) diff --git a/pylabrobot/io/io.py b/pylabrobot/io/io.py index ddd269e1b10..599399a251f 100644 --- a/pylabrobot/io/io.py +++ b/pylabrobot/io/io.py @@ -1,33 +1,8 @@ -import asyncio from abc import ABC, abstractmethod -from typing import TypeVar from pylabrobot.serializer import SerializableMixin -T = TypeVar("T") - - -async def _wait_for_executor_future(future: "asyncio.Future[T]") -> T: - """Wait for executor work to finish, ignoring cancellation of the awaiting task. - - Cancelling the asyncio wrapper returned by ``run_in_executor`` does not stop a callable that is - already running. This keeps teardown from racing work that still owns a device handle. - """ - - async def wait() -> T: - return await future - - waiter = asyncio.create_task(wait()) - while not waiter.done(): - try: - await asyncio.shield(waiter) - except asyncio.CancelledError: - continue - - return waiter.result() - - class IOBase(SerializableMixin, ABC): @abstractmethod async def write(self, data: bytes, *args, **kwargs): diff --git a/pylabrobot/io/serial.py b/pylabrobot/io/serial.py index 590a102416f..6ddc7eaa92d 100644 --- a/pylabrobot/io/serial.py +++ b/pylabrobot/io/serial.py @@ -18,7 +18,6 @@ _SERIAL_IMPORT_ERROR = e from pylabrobot.io.capture import CaptureReader, Command, capturer, get_capture_or_validation_active -from pylabrobot.io.io import _wait_for_executor_future from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences logger = logging.getLogger(__name__) @@ -149,28 +148,23 @@ async def setup(self): self._executor = ThreadPoolExecutor(max_workers=1) setup_future = loop.run_in_executor(self._executor, self._setup_sync) - setup_error: Optional[BaseException] = None try: resolved_port = await asyncio.shield(setup_future) self._port = resolved_port except BaseException as exc: - setup_error = exc - if not setup_future.done(): + if isinstance(exc, asyncio.CancelledError): try: - await _wait_for_executor_future(setup_future) + await setup_future except BaseException: pass - - if setup_error is not None: - if self._ser is not None and self._executor is not None: - close_future = loop.run_in_executor(self._executor, self._ser.close) + if self._ser is not None: try: - await _wait_for_executor_future(close_future) + await loop.run_in_executor(self._executor, self._ser.close) except Exception: logger.warning("Failed to close serial port after setup failure", exc_info=True) - self._ser = None + self._ser = None self._shutdown_executor() - raise setup_error + raise assert self._ser is not None diff --git a/pylabrobot/io/usb.py b/pylabrobot/io/usb.py index fea19116d06..192ad659f87 100644 --- a/pylabrobot/io/usb.py +++ b/pylabrobot/io/usb.py @@ -7,7 +7,7 @@ from pylabrobot.io.capture import Command, capturer, get_capture_or_validation_active from pylabrobot.io.errors import ValidationError -from pylabrobot.io.io import IOBase, _wait_for_executor_future +from pylabrobot.io.io import IOBase from pylabrobot.io.validation_utils import LOG_LEVEL_IO, align_sequences try: @@ -465,32 +465,25 @@ async def setup(self, empty_buffer=True): loop = asyncio.get_running_loop() setup_future = loop.run_in_executor(self._read_executor, self._setup_sync, empty_buffer) - setup_error: Optional[BaseException] = None try: await asyncio.shield(setup_future) except BaseException as exc: - setup_error = exc - if not setup_future.done(): + if isinstance(exc, asyncio.CancelledError): try: - await _wait_for_executor_future(setup_future) + await setup_future except BaseException: pass - - if setup_error is not None: - if self.dev is not None and self._read_executor is not None: + if self.dev is not None: dev = self.dev - dispose_future = loop.run_in_executor( - self._read_executor, lambda: usb.util.dispose_resources(dev) - ) try: - await _wait_for_executor_future(dispose_future) + await loop.run_in_executor(self._read_executor, lambda: usb.util.dispose_resources(dev)) except Exception: logger.warning("Failed to dispose USB device after setup failure", exc_info=True) - self.dev = None - self.read_endpoint = None - self.write_endpoint = None + self.dev = None + self.read_endpoint = None + self.write_endpoint = None self._shutdown_executors() - raise setup_error + raise def _shutdown_executors(self) -> None: for executor in (self._read_executor, self._write_executor):