feat(cloud-scheduler): reserve devices atomically

tests/test_task_scheduler.py

@@ -9,7 +9,6 @@ import pytest
 from cloud.config import CloudConfig
 from cloud.pool import DevicePool
 from cloud.scheduler import (
-    AssignmentStrategy,
     FIFO_MATCH_STRATEGY_NAME,
     QueueFullError,
     ScheduledTask,
@@ -35,11 +34,15 @@ def _config(**overrides) -> CloudConfig:
     return CloudConfig(**base)
 
 
-def _device(device_id: str, *, status: str = "idle", driver_type: str = "wda") -> Device:
+def _device(
+    device_id: str, *, status: str = "idle", driver_type: str = "wda"
+) -> Device:
     return Device(id=device_id, status=status, driver_type=driver_type)  # type: ignore[arg-type]
 
 
-def _pool_with_devices(tmp_path, *devices: Device, host_id: str = "host-local") -> DevicePool:
+def _pool_with_devices(
+    tmp_path, *devices: Device, host_id: str = "host-local"
+) -> DevicePool:
     pool = DevicePool(CloudStore(tmp_path / "cloud.sqlite3"), _config())
     pool.sync_host_devices(host_id, list(devices))
     return pool
@@ -134,6 +137,11 @@ def test_two_tasks_assigned_in_submission_order_with_one_device(tmp_path) -> Non
     assert first_task.status == "assigned"
     assert second_task.status == "queued"
 
+    assert scheduler.assign() == []
+    second_task = pool.store.get_task(second_id)
+    assert second_task is not None
+    assert second_task.status == "queued"
+
 
 def test_unknown_strategy_raises_at_init(tmp_path) -> None:
     pool = _pool_with_devices(tmp_path)