q792602257
1bdc3784dc
- pooled_devices primary key changed from device_id alone to (host_id, device_id), so two hosts reporting the same local device_id no longer crash sync_host_devices() with an uncaught sqlite3.IntegrityError. - Device gained a capability_tags field so DevicePool.sync_host_devices() can actually populate PooledDevice.capability_tags from a real host sync instead of always falling back to an empty list. openspec: device-pool capability, archived change cloud-runtime
157 lines
5.2 KiB
Python
157 lines
5.2 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import replace
|
|
from threading import RLock
|
|
from time import sleep
|
|
|
|
from core.errors import (
|
|
DeviceBusyError,
|
|
DeviceNotFoundError,
|
|
DeviceOfflineError,
|
|
DriverError,
|
|
)
|
|
from core.models import Device, DeviceStatus
|
|
from driver.base import Driver
|
|
|
|
DriverFactory = Callable[[], Driver]
|
|
|
|
|
|
class DeviceManager:
|
|
def __init__(self) -> None:
|
|
self._devices: dict[str, Device] = {}
|
|
self._factories: dict[str, DriverFactory] = {}
|
|
self._drivers: dict[str, Driver] = {}
|
|
self._lock = RLock()
|
|
|
|
def register_device(
|
|
self,
|
|
device_id: str,
|
|
driver_factory: DriverFactory,
|
|
*,
|
|
name: str | None = None,
|
|
driver_type: str = "wda",
|
|
connection_info: dict[str, object] | None = None,
|
|
status: DeviceStatus = "idle",
|
|
capability_tags: list[str] | None = None,
|
|
) -> Device:
|
|
with self._lock:
|
|
device = Device(
|
|
id=device_id,
|
|
name=name,
|
|
status=status,
|
|
driver_type=driver_type,
|
|
connection_info=dict(connection_info or {}),
|
|
capability_tags=list(capability_tags or []),
|
|
)
|
|
self._devices[device_id] = device
|
|
self._factories[device_id] = driver_factory
|
|
self._drivers.pop(device_id, None)
|
|
return replace(device)
|
|
|
|
def unregister_device(self, device_id: str) -> None:
|
|
with self._lock:
|
|
driver = self._drivers.pop(device_id, None)
|
|
if driver:
|
|
driver.disconnect()
|
|
self._devices.pop(device_id, None)
|
|
self._factories.pop(device_id, None)
|
|
|
|
def list_devices(self) -> list[Device]:
|
|
with self._lock:
|
|
return [replace(device) for device in self._devices.values()]
|
|
|
|
def status(self, device_id: str) -> DeviceStatus:
|
|
with self._lock:
|
|
return self._device(device_id).status
|
|
|
|
def connect(
|
|
self,
|
|
device_id: str,
|
|
*,
|
|
max_retries: int = 2,
|
|
retry_backoff_seconds: float = 0.25,
|
|
) -> Driver:
|
|
if max_retries < 1:
|
|
raise ValueError("max_retries must be at least 1")
|
|
|
|
last_error: Exception | None = None
|
|
for attempt in range(1, max_retries + 1):
|
|
with self._lock:
|
|
device = self._device(device_id)
|
|
if device.status == "busy" and device_id in self._drivers:
|
|
return self._drivers[device_id]
|
|
if device.status == "busy":
|
|
raise DeviceBusyError(f"device {device_id} is busy")
|
|
driver = self._factories[device_id]()
|
|
|
|
try:
|
|
driver.connect()
|
|
except Exception as exc:
|
|
last_error = exc
|
|
with self._lock:
|
|
status: DeviceStatus = (
|
|
"offline" if attempt == max_retries else "error"
|
|
)
|
|
self._set_status(device_id, status)
|
|
if attempt < max_retries:
|
|
sleep(retry_backoff_seconds)
|
|
continue
|
|
|
|
with self._lock:
|
|
self._drivers[device_id] = driver
|
|
self._set_status(device_id, "busy")
|
|
return driver
|
|
|
|
raise DeviceOfflineError(
|
|
f"device {device_id} is offline"
|
|
) from last_error
|
|
|
|
def disconnect(self, device_id: str) -> None:
|
|
with self._lock:
|
|
self._device(device_id)
|
|
driver = self._drivers.pop(device_id, None)
|
|
if driver:
|
|
driver.disconnect()
|
|
with self._lock:
|
|
self._set_status(device_id, "idle")
|
|
|
|
def mark_error(self, device_id: str, *, offline: bool = False) -> None:
|
|
with self._lock:
|
|
self._device(device_id)
|
|
self._drivers.pop(device_id, None)
|
|
self._set_status(device_id, "offline" if offline else "error")
|
|
|
|
def active_driver(self, device_id: str | None = None) -> Driver:
|
|
with self._lock:
|
|
if device_id is None:
|
|
busy_ids = [
|
|
known_id
|
|
for known_id, device in self._devices.items()
|
|
if device.status == "busy" and known_id in self._drivers
|
|
]
|
|
if len(busy_ids) != 1:
|
|
raise DriverError(
|
|
"exactly one connected device is required when no device_id is given"
|
|
)
|
|
device_id = busy_ids[0]
|
|
|
|
self._device(device_id)
|
|
driver = self._drivers.get(device_id)
|
|
if not driver:
|
|
raise DeviceOfflineError(f"device {device_id} is not connected")
|
|
return driver
|
|
|
|
def _device(self, device_id: str) -> Device:
|
|
try:
|
|
return self._devices[device_id]
|
|
except KeyError as exc:
|
|
raise DeviceNotFoundError(f"unknown device {device_id}") from exc
|
|
|
|
def _set_status(self, device_id: str, status: DeviceStatus) -> None:
|
|
device = self._devices[device_id]
|
|
self._devices[device_id] = replace(device, status=status)
|
|
|
|
|
|
DEFAULT_MANAGER = DeviceManager()
|