feat(cloud-scheduler): reserve devices atomically

packages/cloud-platform/cloud/scheduler.py

@@ -9,7 +9,7 @@ by ``driver/registry.py``, ``perception/provider.py``, and ``workflow/conditions
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
 from uuid import uuid4
 
@@ -157,32 +157,38 @@ class TaskScheduler:
             return assignments
 
         devices = self.pool.list_devices()
-        assigned_device_ids: set[str] = set()
+        now = utc_now()
+        reserved_device_ids = self.store.list_reserved_device_ids(now=now)
         for task in queued:
             candidates = [
                 device
                 for device in devices
-                if device.device_id not in assigned_device_ids
+                if device.device_id not in reserved_device_ids
                 and device.status == "idle"
                 and _matches(device, task.constraints)
             ]
             selected = strategy.select(task, candidates)
             if selected is None:
                 continue
-            assigned_device_ids.add(selected.device_id)
-            self.store.update_task(
-                task.id,
-                status="assigned",
-                assigned_device_id=selected.device_id,
-                assigned_host_id=selected.host_id,
+            leased = self.store.assign_task(
+                task_id=task.id,
+                host_id=selected.host_id,
+                device_id=selected.device_id,
+                lease_id=uuid4().hex,
+                lease_expires_at=now
+                + timedelta(seconds=self.config.lease_duration_seconds),
+                now=now,
             )
+            if leased is None:
+                continue
+            reserved_device_ids.add(selected.device_id)
             assignments.append(
                 Assignment(
-                    task_id=task.id,
-                    device_id=selected.device_id,
-                    host_id=selected.host_id,
-                    goal=task.goal,
-                    workflow_definition_id=task.workflow_definition_id,
+                    task_id=leased.task_id,
+                    device_id=leased.device_id,
+                    host_id=leased.host_id,
+                    goal=leased.goal,
+                    workflow_definition_id=leased.workflow_definition_id,
                 )
             )
         return assignments