feat(cloud-scheduler): reserve devices atomically

2026-07-12 17:19:10 +08:00
parent 0a9392b7ea
commit ac7734c7dc
6 changed files with 301 additions and 19 deletions
@@ -25,6 +25,9 @@ class CloudConfig:
    # are rejected rather than letting the backlog grow without limit.
    max_queue_depth: int = 100

+    # Duration of each scheduler-created assignment lease.
+    lease_duration_seconds: float = 60.0
+
    # Strategy name looked up in the AssignmentStrategy registry (default fifo_match).
    default_assignment_strategy: str = "fifo_match"

@@ -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
@@ -207,6 +207,89 @@ class SQLAlchemyCloudRepository:
            row = session.get(PluginRow, name)
        return _plugin_from_row(row) if row else None

+    def list_reserved_device_ids(self, *, now: datetime) -> set[str]:
+        with self._sessions() as session:
+            device_ids = session.scalars(
+                select(ScheduledTaskRow.assigned_device_id).where(
+                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                    ScheduledTaskRow.assigned_device_id.is_not(None),
+                    ScheduledTaskRow.lease_expires_at.is_not(None),
+                    ScheduledTaskRow.lease_expires_at > _iso(now),
+                )
+            ).all()
+        return {device_id for device_id in device_ids if device_id is not None}
+
+    def assign_task(
+        self,
+        *,
+        task_id: str,
+        host_id: str,
+        device_id: str,
+        lease_id: str,
+        lease_expires_at: datetime,
+        now: datetime,
+    ) -> Any | None:
+        with self._sessions.begin() as session:
+            task_statement = select(ScheduledTaskRow).where(
+                ScheduledTaskRow.id == task_id
+            )
+            device_statement = select(PooledDeviceRow).where(
+                PooledDeviceRow.host_id == host_id,
+                PooledDeviceRow.device_id == device_id,
+            )
+            if self.engine.dialect.name == "postgresql":
+                task_statement = task_statement.with_for_update(skip_locked=True)
+                device_statement = device_statement.with_for_update()
+
+            task = session.scalars(task_statement).first()
+            if task is None or task.status != "queued":
+                return None
+            device = session.scalars(device_statement).first()
+            if device is None or device.status != "idle":
+                return None
+
+            active_reservation = session.scalar(
+                select(ScheduledTaskRow.id)
+                .where(
+                    ScheduledTaskRow.id != task_id,
+                    ScheduledTaskRow.assigned_device_id == device_id,
+                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                    ScheduledTaskRow.lease_expires_at.is_not(None),
+                    ScheduledTaskRow.lease_expires_at > _iso(now),
+                )
+                .limit(1)
+            )
+            if active_reservation is not None:
+                return None
+
+            attempt = task.attempt_count + 1
+            task.status = "assigned"
+            task.assigned_host_id = host_id
+            task.assigned_device_id = device_id
+            task.attempt_count = attempt
+            task.lease_id = lease_id
+            task.lease_expires_at = _iso(lease_expires_at)
+            task.failure_reason = None
+            task.result_json = None
+            task.updated_at = _iso(now)
+            session.add(
+                TaskAttemptRow(
+                    task_id=task.id,
+                    attempt=attempt,
+                    lease_id=lease_id,
+                    host_id=host_id,
+                    device_id=device_id,
+                    status="assigned",
+                    lease_expires_at=_iso(lease_expires_at),
+                    created_at=_iso(now),
+                    completed_at=None,
+                    failure_reason=None,
+                    result_json=None,
+                )
+            )
+            session.flush()
+            return _leased_assignment_from_row(task)
+
    def list_task_attempts(self, task_id: str) -> list[Any]:
        with self._sessions() as session:
            rows = session.scalars(
@@ -311,6 +394,29 @@ def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
    )


+def _leased_assignment_from_row(row: ScheduledTaskRow) -> Any:
+    from cloud.repository import LeasedAssignment
+
+    lease_expires_at = _parse_dt(row.lease_expires_at)
+    if (
+        row.lease_id is None
+        or lease_expires_at is None
+        or row.assigned_host_id is None
+        or row.assigned_device_id is None
+    ):
+        raise ValueError(f"task {row.id!r} does not contain a complete lease")
+    return LeasedAssignment(
+        task_id=row.id,
+        attempt=row.attempt_count,
+        lease_id=row.lease_id,
+        lease_expires_at=lease_expires_at,
+        host_id=row.assigned_host_id,
+        device_id=row.assigned_device_id,
+        goal=row.goal,
+        workflow_definition_id=row.workflow_definition_id,
+    )
+
+
 def _parse_json_object(value: str | None) -> dict[str, Any] | None:
    if value is None:
        return None