From ef4c1bd3927b786446ec495056ef586a02311f6d Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 20 Aug 2025 17:06:44 -0400 Subject: [PATCH 1/7] treewide: update anyio to >=4.10.0 --- __templates__/driver/pyproject.toml.tmpl | 2 +- .../pyproject.toml | 2 +- .../pyproject.toml | 2 +- .../pyproject.toml | 2 +- .../jumpstarter-driver-http/pyproject.toml | 2 +- .../jumpstarter-driver-iscsi/pyproject.toml | 2 +- .../pyproject.toml | 2 +- .../jumpstarter-driver-shell/pyproject.toml | 2 +- .../jumpstarter-driver-tasmota/pyproject.toml | 2 +- .../jumpstarter-driver-tftp/pyproject.toml | 2 +- .../jumpstarter-driver-yepkit/pyproject.toml | 2 +- uv.lock | 92 +++++++++---------- 12 files changed, 57 insertions(+), 57 deletions(-) diff --git a/__templates__/driver/pyproject.toml.tmpl b/__templates__/driver/pyproject.toml.tmpl index e7694b7c7..e7dbb11ee 100644 --- a/__templates__/driver/pyproject.toml.tmpl +++ b/__templates__/driver/pyproject.toml.tmpl @@ -9,7 +9,7 @@ authors = [ ] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter", ] diff --git a/packages/jumpstarter-driver-energenie/pyproject.toml b/packages/jumpstarter-driver-energenie/pyproject.toml index bfe5c6a46..0e5f22e1c 100644 --- a/packages/jumpstarter-driver-energenie/pyproject.toml +++ b/packages/jumpstarter-driver-energenie/pyproject.toml @@ -9,7 +9,7 @@ authors = [ ] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter", "jumpstarter-driver-power" ] diff --git a/packages/jumpstarter-driver-flashers/pyproject.toml b/packages/jumpstarter-driver-flashers/pyproject.toml index d026c535b..2127a6fae 100644 --- a/packages/jumpstarter-driver-flashers/pyproject.toml +++ b/packages/jumpstarter-driver-flashers/pyproject.toml @@ -11,7 +11,7 @@ authors = [ requires-python = ">=3.11" dependencies = [ "oras>=0.2.25", - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter", "jumpstarter-driver-opendal", "jumpstarter-driver-pyserial", diff --git a/packages/jumpstarter-driver-http-power/pyproject.toml b/packages/jumpstarter-driver-http-power/pyproject.toml index d6549ac9d..5791f9a25 100644 --- a/packages/jumpstarter-driver-http-power/pyproject.toml +++ b/packages/jumpstarter-driver-http-power/pyproject.toml @@ -9,7 +9,7 @@ authors = [ ] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter", "jumpstarter-driver-power", ] diff --git a/packages/jumpstarter-driver-http/pyproject.toml b/packages/jumpstarter-driver-http/pyproject.toml index db47e88d7..55d6053f6 100644 --- a/packages/jumpstarter-driver-http/pyproject.toml +++ b/packages/jumpstarter-driver-http/pyproject.toml @@ -7,7 +7,7 @@ license = "Apache-2.0" authors = [{ name = "Benny Zlotnik", email = "bzlotnik@redhat.com" }] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter", "jumpstarter-driver-composite", "jumpstarter-driver-opendal", diff --git a/packages/jumpstarter-driver-iscsi/pyproject.toml b/packages/jumpstarter-driver-iscsi/pyproject.toml index 8259bcfe7..f5569065d 100644 --- a/packages/jumpstarter-driver-iscsi/pyproject.toml +++ b/packages/jumpstarter-driver-iscsi/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" authors = [{ name = "Benny Zlotnik", email = "bzlotnik@redhat.com" }] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter", "jumpstarter-driver-composite", "jumpstarter-driver-opendal", diff --git a/packages/jumpstarter-driver-probe-rs/pyproject.toml b/packages/jumpstarter-driver-probe-rs/pyproject.toml index d1880e411..16243c0bb 100644 --- a/packages/jumpstarter-driver-probe-rs/pyproject.toml +++ b/packages/jumpstarter-driver-probe-rs/pyproject.toml @@ -7,7 +7,7 @@ license = "Apache-2.0" authors = [{ name = "Miguel Angel Ajo", email = "miguelangel@ajo.es" }] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "click>=8.1.7.2", "jumpstarter", "jumpstarter-driver-opendal", diff --git a/packages/jumpstarter-driver-shell/pyproject.toml b/packages/jumpstarter-driver-shell/pyproject.toml index 670da0728..08b642387 100644 --- a/packages/jumpstarter-driver-shell/pyproject.toml +++ b/packages/jumpstarter-driver-shell/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" authors = [{ name = "Miguel Angel Ajo", email = "miguelangel@ajo.es" }] requires-python = ">=3.11" license = "Apache-2.0" -dependencies = ["anyio>=4.6.2.post1", "jumpstarter"] +dependencies = ["anyio>=4.10.0", "jumpstarter"] [project.entry-points."jumpstarter.drivers"] Shell = "jumpstarter_driver_shell.driver:Shell" diff --git a/packages/jumpstarter-driver-tasmota/pyproject.toml b/packages/jumpstarter-driver-tasmota/pyproject.toml index a8eda3dd8..ae32a1193 100644 --- a/packages/jumpstarter-driver-tasmota/pyproject.toml +++ b/packages/jumpstarter-driver-tasmota/pyproject.toml @@ -7,7 +7,7 @@ license = "Apache-2.0" authors = [{ name = "Nick Cao", email = "nickcao@nichi.co" }] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter_driver_power", "jumpstarter", "paho-mqtt>=2.1.0", diff --git a/packages/jumpstarter-driver-tftp/pyproject.toml b/packages/jumpstarter-driver-tftp/pyproject.toml index 2debc7196..bb3429a7b 100644 --- a/packages/jumpstarter-driver-tftp/pyproject.toml +++ b/packages/jumpstarter-driver-tftp/pyproject.toml @@ -7,7 +7,7 @@ license = "Apache-2.0" authors = [{ name = "Benny Zlotnik", email = "bzlotnik@redhat.com" }] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "jumpstarter", "jumpstarter-driver-composite", "jumpstarter-driver-opendal", diff --git a/packages/jumpstarter-driver-yepkit/pyproject.toml b/packages/jumpstarter-driver-yepkit/pyproject.toml index a851197f9..13607a5d9 100644 --- a/packages/jumpstarter-driver-yepkit/pyproject.toml +++ b/packages/jumpstarter-driver-yepkit/pyproject.toml @@ -7,7 +7,7 @@ license = "Apache-2.0" authors = [{ name = "Miguel Angel Ajo", email = "miguelangel@ajo.es" }] requires-python = ">=3.11" dependencies = [ - "anyio>=4.6.2.post1", + "anyio>=4.10.0", "pyusb>=1.2.1", "jumpstarter_driver_power", "jumpstarter", diff --git a/uv.lock b/uv.lock index 8f795e1a5..ed2d7e723 100644 --- a/uv.lock +++ b/uv.lock @@ -17,6 +17,7 @@ members = [ "jumpstarter-driver-dutlink", "jumpstarter-driver-energenie", "jumpstarter-driver-flashers", + "jumpstarter-driver-gpiod", "jumpstarter-driver-http", "jumpstarter-driver-http-power", "jumpstarter-driver-iscsi", @@ -26,7 +27,6 @@ members = [ "jumpstarter-driver-probe-rs", "jumpstarter-driver-pyserial", "jumpstarter-driver-qemu", - "jumpstarter-driver-gpiod", "jumpstarter-driver-ridesx", "jumpstarter-driver-sdwire", "jumpstarter-driver-shell", @@ -176,16 +176,16 @@ wheels = [ [[package]] name = "anyio" -version = "4.9.0" +version = "4.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "idna" }, { name = "sniffio" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/95/7d/4c1bd541d4dffa1b52bd83fb8527089e097a106fc90b467a7313b105f840/anyio-4.9.0.tar.gz", hash = "sha256:673c0c244e15788651a4ff38710fea9675823028a6f08a5eda409e0c9840a028", size = 190949, upload-time = "2025-03-17T00:02:54.77Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/b4/636b3b65173d3ce9a38ef5f0522789614e590dab6a8d505340a4efe4c567/anyio-4.10.0.tar.gz", hash = "sha256:3f3fae35c96039744587aa5b8371e7e8e603c0702999535961dd336026973ba6", size = 213252, upload-time = "2025-08-04T08:54:26.451Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a1/ee/48ca1a7c89ffec8b6a0c5d02b89c305671d5ffd8d3c94acf8b8c408575bb/anyio-4.9.0-py3-none-any.whl", hash = "sha256:9f76d541cad6e36af7beb62e978876f3b41e3e04f2c1fbf0884604c0a9c4d93c", size = 100916, upload-time = "2025-03-17T00:02:52.713Z" }, + { url = "https://files.pythonhosted.org/packages/6f/12/e5e0282d673bb9746bacfb6e2dba8719989d3660cdb2ea79aee9a9651afb/anyio-4.10.0-py3-none-any.whl", hash = "sha256:60e474ac86736bbfd6f210f7a61218939c318f43f9972497381f1c5e930ed3d1", size = 107213, upload-time = "2025-08-04T08:54:24.882Z" }, ] [[package]] @@ -1064,6 +1064,7 @@ dependencies = [ { name = "jumpstarter-driver-corellium" }, { name = "jumpstarter-driver-dutlink" }, { name = "jumpstarter-driver-flashers" }, + { name = "jumpstarter-driver-gpiod" }, { name = "jumpstarter-driver-http" }, { name = "jumpstarter-driver-network" }, { name = "jumpstarter-driver-opendal" }, @@ -1071,7 +1072,6 @@ dependencies = [ { name = "jumpstarter-driver-probe-rs" }, { name = "jumpstarter-driver-pyserial" }, { name = "jumpstarter-driver-qemu" }, - { name = "jumpstarter-driver-gpiod" }, { name = "jumpstarter-driver-ridesx" }, { name = "jumpstarter-driver-sdwire" }, { name = "jumpstarter-driver-shell" }, @@ -1098,6 +1098,7 @@ requires-dist = [ { name = "jumpstarter-driver-corellium", editable = "packages/jumpstarter-driver-corellium" }, { name = "jumpstarter-driver-dutlink", editable = "packages/jumpstarter-driver-dutlink" }, { name = "jumpstarter-driver-flashers", editable = "packages/jumpstarter-driver-flashers" }, + { name = "jumpstarter-driver-gpiod", editable = "packages/jumpstarter-driver-gpiod" }, { name = "jumpstarter-driver-http", editable = "packages/jumpstarter-driver-http" }, { name = "jumpstarter-driver-network", editable = "packages/jumpstarter-driver-network" }, { name = "jumpstarter-driver-opendal", editable = "packages/jumpstarter-driver-opendal" }, @@ -1105,7 +1106,6 @@ requires-dist = [ { name = "jumpstarter-driver-probe-rs", editable = "packages/jumpstarter-driver-probe-rs" }, { name = "jumpstarter-driver-pyserial", editable = "packages/jumpstarter-driver-pyserial" }, { name = "jumpstarter-driver-qemu", editable = "packages/jumpstarter-driver-qemu" }, - { name = "jumpstarter-driver-gpiod", editable = "packages/jumpstarter-driver-gpiod" }, { name = "jumpstarter-driver-ridesx", editable = "packages/jumpstarter-driver-ridesx" }, { name = "jumpstarter-driver-sdwire", editable = "packages/jumpstarter-driver-sdwire" }, { name = "jumpstarter-driver-shell", editable = "packages/jumpstarter-driver-shell" }, @@ -1412,7 +1412,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-power", editable = "packages/jumpstarter-driver-power" }, ] @@ -1447,7 +1447,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-http", editable = "packages/jumpstarter-driver-http" }, { name = "jumpstarter-driver-opendal", editable = "packages/jumpstarter-driver-opendal" }, @@ -1464,6 +1464,36 @@ dev = [ { name = "pytest-cov", specifier = ">=6.0.0" }, ] +[[package]] +name = "jumpstarter-driver-gpiod" +source = { editable = "packages/jumpstarter-driver-gpiod" } +dependencies = [ + { name = "click" }, + { name = "gpiod", marker = "sys_platform == 'linux'" }, + { name = "jumpstarter" }, + { name = "jumpstarter-driver-power" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pytest" }, + { name = "pytest-cov" }, +] + +[package.metadata] +requires-dist = [ + { name = "click", specifier = ">=8.1.7.2" }, + { name = "gpiod", marker = "sys_platform == 'linux'" }, + { name = "jumpstarter", editable = "packages/jumpstarter" }, + { name = "jumpstarter-driver-power", editable = "packages/jumpstarter-driver-power" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pytest", specifier = ">=8.3.2" }, + { name = "pytest-cov", specifier = ">=5.0.0" }, +] + [[package]] name = "jumpstarter-driver-http" source = { editable = "packages/jumpstarter-driver-http" } @@ -1484,7 +1514,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-composite", editable = "packages/jumpstarter-driver-composite" }, { name = "jumpstarter-driver-opendal", editable = "packages/jumpstarter-driver-opendal" }, @@ -1516,7 +1546,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-power", editable = "packages/jumpstarter-driver-power" }, ] @@ -1547,7 +1577,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-composite", editable = "packages/jumpstarter-driver-composite" }, { name = "jumpstarter-driver-opendal", editable = "packages/jumpstarter-driver-opendal" }, @@ -1679,7 +1709,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "click", specifier = ">=8.1.7.2" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-opendal", editable = "packages/jumpstarter-driver-opendal" }, @@ -1767,36 +1797,6 @@ dev = [ { name = "requests", specifier = ">=2.32.3" }, ] -[[package]] -name = "jumpstarter-driver-gpiod" -source = { editable = "packages/jumpstarter-driver-gpiod" } -dependencies = [ - { name = "click" }, - { name = "gpiod", marker = "sys_platform == 'linux'" }, - { name = "jumpstarter" }, - { name = "jumpstarter-driver-power" }, -] - -[package.dev-dependencies] -dev = [ - { name = "pytest" }, - { name = "pytest-cov" }, -] - -[package.metadata] -requires-dist = [ - { name = "click", specifier = ">=8.1.7.2" }, - { name = "gpiod", marker = "sys_platform == 'linux'" }, - { name = "jumpstarter", editable = "packages/jumpstarter" }, - { name = "jumpstarter-driver-power", editable = "packages/jumpstarter-driver-power" }, -] - -[package.metadata.requires-dev] -dev = [ - { name = "pytest", specifier = ">=8.3.2" }, - { name = "pytest-cov", specifier = ">=5.0.0" }, -] - [[package]] name = "jumpstarter-driver-ridesx" source = { editable = "packages/jumpstarter-driver-ridesx" } @@ -1871,7 +1871,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, ] @@ -1936,7 +1936,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-power", editable = "packages/jumpstarter-driver-power" }, { name = "paho-mqtt", specifier = ">=2.1.0" }, @@ -1970,7 +1970,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-composite", editable = "packages/jumpstarter-driver-composite" }, { name = "jumpstarter-driver-opendal", editable = "packages/jumpstarter-driver-opendal" }, @@ -2065,7 +2065,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "anyio", specifier = ">=4.6.2.post1" }, + { name = "anyio", specifier = ">=4.10.0" }, { name = "jumpstarter", editable = "packages/jumpstarter" }, { name = "jumpstarter-driver-power", editable = "packages/jumpstarter-driver-power" }, { name = "pyusb", specifier = ">=1.2.1" }, From edc75de5c4cd4976340ac3bd1db0b4ec0f6548fd Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 20 Aug 2025 17:14:37 -0400 Subject: [PATCH 2/7] Lease: use ContextManagerMixin --- .../jumpstarter/jumpstarter/client/lease.py | 28 ++++++++----------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/packages/jumpstarter/jumpstarter/client/lease.py b/packages/jumpstarter/jumpstarter/client/lease.py index f9fc8bea3..fdbc1f8b4 100644 --- a/packages/jumpstarter/jumpstarter/client/lease.py +++ b/packages/jumpstarter/jumpstarter/client/lease.py @@ -1,16 +1,15 @@ import logging +from collections.abc import AsyncGenerator, Generator from contextlib import ( - AbstractAsyncContextManager, - AbstractContextManager, ExitStack, asynccontextmanager, contextmanager, ) from dataclasses import dataclass, field from datetime import datetime, timedelta -from typing import Any +from typing import Any, Self -from anyio import create_task_group, fail_after, sleep +from anyio import AsyncContextManagerMixin, ContextManagerMixin, create_task_group, fail_after, sleep from anyio.from_thread import BlockingPortal from grpc.aio import Channel from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc @@ -28,7 +27,7 @@ @dataclass(kw_only=True) -class Lease(AbstractContextManager, AbstractAsyncContextManager): +class Lease(ContextManagerMixin, AsyncContextManagerMixin): channel: Channel duration: timedelta selector: str @@ -48,7 +47,6 @@ def __post_init__(self): self.controller = jumpstarter_pb2_grpc.ControllerServiceStub(self.channel) self.svc = ClientService(channel=self.channel, namespace=self.namespace) - self.manager = self.portal.wrap_async_context_manager(self) async def _create(self): logger.debug("Creating lease request for selector %s for duration %s", self.selector, self.duration) @@ -127,23 +125,19 @@ async def _acquire(self): await sleep(1) - async def __aenter__(self): - return await self.request_async() - - async def __aexit__(self, exc_type, exc_value, traceback): + @asynccontextmanager + async def __asynccontextmanager__(self) -> AsyncGenerator[Self]: + yield await self.request_async() if self.release: logger.info("Releasing Lease %s", self.name) await self.svc.DeleteLease( name=self.name, ) - def __enter__(self): - # wraps the async context manager enter - return self.manager.__enter__() - - def __exit__(self, exc_type, exc_value, traceback): - # wraps the async context manager exit - return self.manager.__exit__(exc_type, exc_value, traceback) + @contextmanager + def __contextmanager__(self) -> Generator[Self]: + with self.portal.wrap_async_context_manager(self) as value: + yield value async def handle_async(self, stream): logger.debug("Connecting to Lease with name %s", self.name) From 40f9e87342c96dc1d54905d4550883c732461c7f Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 20 Aug 2025 17:17:11 -0400 Subject: [PATCH 3/7] Session: use ContextManagerMixin --- .../jumpstarter/jumpstarter/exporter/session.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/packages/jumpstarter/jumpstarter/exporter/session.py b/packages/jumpstarter/jumpstarter/exporter/session.py index 7334e235c..8655b1474 100644 --- a/packages/jumpstarter/jumpstarter/exporter/session.py +++ b/packages/jumpstarter/jumpstarter/exporter/session.py @@ -1,12 +1,14 @@ import logging from collections import deque -from contextlib import AbstractContextManager, asynccontextmanager, contextmanager, suppress +from collections.abc import Generator +from contextlib import asynccontextmanager, contextmanager, suppress from dataclasses import dataclass, field from logging.handlers import QueueHandler +from typing import Self from uuid import UUID import grpc -from anyio import Event, TypedAttributeLookupError, sleep +from anyio import ContextManagerMixin, Event, TypedAttributeLookupError, sleep from anyio.from_thread import start_blocking_portal from jumpstarter_protocol import ( jumpstarter_pb2, @@ -30,7 +32,7 @@ class Session( jumpstarter_pb2_grpc.ExporterServiceServicer, router_pb2_grpc.RouterServiceServicer, Metadata, - AbstractContextManager, + ContextManagerMixin, ): root_device: Driver mapping: dict[UUID, Driver] @@ -38,19 +40,18 @@ class Session( _logging_queue: deque = field(init=False) _logging_handler: QueueHandler = field(init=False) - def __enter__(self): + @contextmanager + def __contextmanager__(self) -> Generator[Self]: logging.getLogger().addHandler(self._logging_handler) self.root_device.reset() - return self - - def __exit__(self, exc_type, exc_value, traceback): + yield self try: self.root_device.close() except Exception as e: # Get driver name from report for more descriptive logging try: report = self.root_device.report() - driver_name = report.labels.get('jumpstarter.dev/name', self.root_device.__class__.__name__) + driver_name = report.labels.get("jumpstarter.dev/name", self.root_device.__class__.__name__) except Exception: driver_name = self.root_device.__class__.__name__ logger.error("Error closing driver %s: %s", driver_name, e, exc_info=True) From 2f5123b810feefddf06e7b537709b1fe14100d38 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 20 Aug 2025 17:19:19 -0400 Subject: [PATCH 4/7] DbusNetworkClient: use ContextManagerMixin --- .../jumpstarter_driver_network/client.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py b/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py index 2b18e8d7e..4d2ad05f6 100644 --- a/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py +++ b/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py @@ -1,8 +1,11 @@ -from contextlib import AbstractContextManager +from collections.abc import Generator +from contextlib import contextmanager from ipaddress import IPv6Address, ip_address from threading import Event +from typing import Any import click +from anyio import ContextManagerMixin from .adapters import DbusAdapter, TcpPortforwardAdapter, UnixPortforwardAdapter from .driver import DbusNetwork @@ -62,13 +65,11 @@ def forward_unix(path: str | None): return base -class DbusNetworkClient(NetworkClient, AbstractContextManager): - def __enter__(self): - self.adapter = DbusAdapter(client=self) - self.adapter.__enter__() - - def __exit__(self, exc_type, exc_value, traceback): - self.adapter.__exit__(exc_type, exc_value, traceback) +class DbusNetworkClient(NetworkClient, ContextManagerMixin): + @contextmanager + def __contextmanager__(self) -> Generator[Any]: + with DbusAdapter(client=self) as value: + yield value @property def kind(self): From 2049a0cc8fc9db7a39a5c70450a055cefbf1b9e7 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 20 Aug 2025 17:24:59 -0400 Subject: [PATCH 5/7] Exporter: use ContextManagerMixin --- .../jumpstarter/exporter/exporter.py | 67 +++++++++++-------- 1 file changed, 39 insertions(+), 28 deletions(-) diff --git a/packages/jumpstarter/jumpstarter/exporter/exporter.py b/packages/jumpstarter/jumpstarter/exporter/exporter.py index 8fda8fbe2..9fe251e13 100644 --- a/packages/jumpstarter/jumpstarter/exporter/exporter.py +++ b/packages/jumpstarter/jumpstarter/exporter/exporter.py @@ -1,10 +1,19 @@ import logging -from collections.abc import Awaitable, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from collections.abc import AsyncGenerator, Awaitable, Callable +from contextlib import asynccontextmanager from dataclasses import dataclass, field +from typing import Self import grpc -from anyio import connect_unix, create_memory_object_stream, create_task_group, sleep +from anyio import ( + AsyncContextManagerMixin, + CancelScope, + connect_unix, + create_memory_object_stream, + create_task_group, + sleep, + move_on_after, +) from anyio.abc import TaskGroup from google.protobuf import empty_pb2 from jumpstarter_protocol import ( @@ -22,7 +31,7 @@ @dataclass(kw_only=True) -class Exporter(AbstractAsyncContextManager, Metadata): +class Exporter(AsyncContextManagerMixin, Metadata): channel_factory: Callable[[], Awaitable[grpc.aio.Channel]] device_factory: Callable[[], Driver] lease_name: str = field(init=False, default="") @@ -48,32 +57,34 @@ def stop(self, wait_for_lease_exit=False): self._stop_requested = True logger.info("Exporter marked for stop upon lease exit") - async def __aexit__(self, exc_type, exc_value, traceback): - import anyio - + @asynccontextmanager + async def __asynccontextmanager__(self) -> AsyncGenerator[Self]: try: - if self.registered: - logger.info("Unregistering exporter with controller") - try: - with anyio.move_on_after(10): # 10 second timeout - channel = await self.channel_factory() - try: - controller = jumpstarter_pb2_grpc.ControllerServiceStub(channel) - await controller.Unregister( - jumpstarter_pb2.UnregisterRequest( - reason="Exporter shutdown", + yield self + finally: + try: + if self.registered: + logger.info("Unregistering exporter with controller") + try: + with move_on_after(10): # 10 second timeout + channel = await self.channel_factory() + try: + controller = jumpstarter_pb2_grpc.ControllerServiceStub(channel) + await controller.Unregister( + jumpstarter_pb2.UnregisterRequest( + reason="Exporter shutdown", + ) ) - ) - logger.info("Controller unregistration completed successfully") - finally: - with anyio.CancelScope(shield=True): - await channel.close() - except Exception as e: - logger.error("Error during controller unregistration: %s", e, exc_info=True) - - except Exception as e: - logger.error("Error during exporter cleanup: %s", e, exc_info=True) - # Don't re-raise to avoid masking the original exception + logger.info("Controller unregistration completed successfully") + finally: + with CancelScope(shield=True): + await channel.close() + except Exception as e: + logger.error("Error during controller unregistration: %s", e, exc_info=True) + + except Exception as e: + logger.error("Error during exporter cleanup: %s", e, exc_info=True) + # Don't re-raise to avoid masking the original exception async def __handle(self, path, endpoint, token, tls_config, grpc_options): try: From 12cc998804f47530d82c847636981c81d40425ab Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 3 Sep 2025 13:48:01 -0400 Subject: [PATCH 6/7] Fixup error handling in context managers --- .../jumpstarter/jumpstarter/client/lease.py | 15 +++++++----- .../jumpstarter/exporter/session.py | 24 ++++++++++--------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/packages/jumpstarter/jumpstarter/client/lease.py b/packages/jumpstarter/jumpstarter/client/lease.py index fdbc1f8b4..f63e5b437 100644 --- a/packages/jumpstarter/jumpstarter/client/lease.py +++ b/packages/jumpstarter/jumpstarter/client/lease.py @@ -127,12 +127,15 @@ async def _acquire(self): @asynccontextmanager async def __asynccontextmanager__(self) -> AsyncGenerator[Self]: - yield await self.request_async() - if self.release: - logger.info("Releasing Lease %s", self.name) - await self.svc.DeleteLease( - name=self.name, - ) + value = await self.request_async() + try: + yield value + finally: + if self.release: + logger.info("Releasing Lease %s", self.name) + await self.svc.DeleteLease( + name=self.name, + ) @contextmanager def __contextmanager__(self) -> Generator[Self]: diff --git a/packages/jumpstarter/jumpstarter/exporter/session.py b/packages/jumpstarter/jumpstarter/exporter/session.py index 8655b1474..63ae2f08d 100644 --- a/packages/jumpstarter/jumpstarter/exporter/session.py +++ b/packages/jumpstarter/jumpstarter/exporter/session.py @@ -44,19 +44,21 @@ class Session( def __contextmanager__(self) -> Generator[Self]: logging.getLogger().addHandler(self._logging_handler) self.root_device.reset() - yield self try: - self.root_device.close() - except Exception as e: - # Get driver name from report for more descriptive logging - try: - report = self.root_device.report() - driver_name = report.labels.get("jumpstarter.dev/name", self.root_device.__class__.__name__) - except Exception: - driver_name = self.root_device.__class__.__name__ - logger.error("Error closing driver %s: %s", driver_name, e, exc_info=True) + yield self finally: - logging.getLogger().removeHandler(self._logging_handler) + try: + self.root_device.close() + except Exception as e: + # Get driver name from report for more descriptive logging + try: + report = self.root_device.report() + driver_name = report.labels.get("jumpstarter.dev/name", self.root_device.__class__.__name__) + except Exception: + driver_name = self.root_device.__class__.__name__ + logger.error("Error closing driver %s: %s", driver_name, e, exc_info=True) + finally: + logging.getLogger().removeHandler(self._logging_handler) def __init__(self, *args, root_device, **kwargs): super().__init__(*args, **kwargs) From 9fa5a8d23f8bde61ee3b30653f353fb44be89baf Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 3 Sep 2025 13:49:25 -0400 Subject: [PATCH 7/7] fixup! Fixup error handling in context managers --- packages/jumpstarter/jumpstarter/exporter/exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/jumpstarter/jumpstarter/exporter/exporter.py b/packages/jumpstarter/jumpstarter/exporter/exporter.py index 9fe251e13..e63589ee5 100644 --- a/packages/jumpstarter/jumpstarter/exporter/exporter.py +++ b/packages/jumpstarter/jumpstarter/exporter/exporter.py @@ -11,8 +11,8 @@ connect_unix, create_memory_object_stream, create_task_group, - sleep, move_on_after, + sleep, ) from anyio.abc import TaskGroup from google.protobuf import empty_pb2