From ac7734c7dc9886a31d8a92848455e808d471e737 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 17:19:10 +0800 Subject: [PATCH] feat(cloud-scheduler): reserve devices atomically --- .../cloud-control-plane-integration/tasks.md | 2 +- packages/cloud-platform/cloud/config.py | 3 + packages/cloud-platform/cloud/scheduler.py | 34 ++-- .../cloud-platform/cloud/sql_repository.py | 106 ++++++++++++ tests/test_cloud_repository_contract.py | 161 +++++++++++++++++- tests/test_task_scheduler.py | 14 +- 6 files changed, 301 insertions(+), 19 deletions(-) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index a526e3a..6301c59 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -18,7 +18,7 @@ ## 3. Lease-Backed Scheduling - [x] 3.1 Extend scheduled-task persistence with attempt count, lease id/expiry, terminal result, failure reason, and auditable attempt records. -- [ ] 3.2 Implement atomic queued-task assignment and device reservation while excluding devices with active assignments even when snapshots report idle. +- [x] 3.2 Implement atomic queued-task assignment and device reservation while excluding devices with active assignments even when snapshots report idle. - [ ] 3.3 Implement owning-host claim that atomically transitions one assigned attempt to dispatched under its active lease. - [ ] 3.4 Implement lease renewal with host/task/attempt ownership validation and conflict responses for stale leases. - [ ] 3.5 Implement idempotent terminal result recording and reservation release for active leases. diff --git a/packages/cloud-platform/cloud/config.py b/packages/cloud-platform/cloud/config.py index deff5fc..273b59b 100644 --- a/packages/cloud-platform/cloud/config.py +++ b/packages/cloud-platform/cloud/config.py @@ -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" diff --git a/packages/cloud-platform/cloud/scheduler.py b/packages/cloud-platform/cloud/scheduler.py index be4d63d..e687281 100644 --- a/packages/cloud-platform/cloud/scheduler.py +++ b/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 diff --git a/packages/cloud-platform/cloud/sql_repository.py b/packages/cloud-platform/cloud/sql_repository.py index 66b96cd..7b775e4 100644 --- a/packages/cloud-platform/cloud/sql_repository.py +++ b/packages/cloud-platform/cloud/sql_repository.py @@ -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 diff --git a/tests/test_cloud_repository_contract.py b/tests/test_cloud_repository_contract.py index 09de4bc..cbbae36 100644 --- a/tests/test_cloud_repository_contract.py +++ b/tests/test_cloud_repository_contract.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from pathlib import Path from typing import get_protocol_members from uuid import uuid4 @@ -278,3 +278,162 @@ def test_task_attempt_history_is_ordered_and_complete(database_url: str) -> None assert attempts[1].failure_reason == "execution failed" finally: database.close() + + +def test_atomic_assignment_creates_lease_attempt_and_reservation( + database_url: str, +) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("assignment-host") + device_id = _unique_id("assignment-device") + task_id = _unique_id("assignment-task") + now = datetime(2026, 7, 12, 2, 0, tzinfo=UTC) + lease_expires_at = now + timedelta(minutes=1) + + try: + database.repository.upsert_host(host_id, address=None, last_seen_at=now) + database.repository.replace_host_devices( + host_id, + [_device(device_id, host_id)], + ) + database.repository.enqueue_task( + ScheduledTask( + id=task_id, + goal="open settings", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=now, + ) + ) + + assignment = database.repository.assign_task( + task_id=task_id, + host_id=host_id, + device_id=device_id, + lease_id="lease-1", + lease_expires_at=lease_expires_at, + now=now, + ) + + assert assignment is not None + assert assignment.attempt == 1 + assert assignment.lease_id == "lease-1" + task = database.repository.get_task(task_id) + assert task is not None + assert task.status == "assigned" + assert task.attempt_count == 1 + assert task.lease_expires_at == lease_expires_at + assert device_id in database.repository.list_reserved_device_ids(now=now) + attempts = database.repository.list_task_attempts(task_id) + assert len(attempts) == 1 + assert attempts[0].status == "assigned" + assert attempts[0].lease_id == "lease-1" + finally: + database.close() + + +def test_active_assignment_blocks_reuse_of_stale_idle_snapshot( + database_url: str, +) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("reservation-host") + device_id = _unique_id("reservation-device") + first_task_id = _unique_id("reservation-task") + second_task_id = _unique_id("reservation-task") + now = datetime(2026, 7, 12, 3, 0, tzinfo=UTC) + + try: + database.repository.upsert_host(host_id, address=None, last_seen_at=now) + database.repository.replace_host_devices( + host_id, + [_device(device_id, host_id)], + ) + for task_id in (first_task_id, second_task_id): + database.repository.enqueue_task( + ScheduledTask( + id=task_id, + goal="run task", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=now, + ) + ) + + first = database.repository.assign_task( + task_id=first_task_id, + host_id=host_id, + device_id=device_id, + lease_id="lease-active", + lease_expires_at=now + timedelta(minutes=1), + now=now, + ) + second = database.repository.assign_task( + task_id=second_task_id, + host_id=host_id, + device_id=device_id, + lease_id="lease-blocked", + lease_expires_at=now + timedelta(minutes=1), + now=now, + ) + + assert first is not None + assert second is None + blocked_task = database.repository.get_task(second_task_id) + assert blocked_task is not None + assert blocked_task.status == "queued" + assert database.repository.list_task_attempts(second_task_id) == [] + finally: + database.close() + + +def test_expired_assignment_no_longer_reserves_device(database_url: str) -> None: + database = CloudDatabase(database_url) + host_id = _unique_id("expired-host") + device_id = _unique_id("expired-device") + first_task_id = _unique_id("expired-task") + second_task_id = _unique_id("expired-task") + assigned_at = datetime(2026, 7, 12, 4, 0, tzinfo=UTC) + after_expiry = assigned_at + timedelta(minutes=2) + + try: + database.repository.upsert_host( + host_id, + address=None, + last_seen_at=assigned_at, + ) + database.repository.replace_host_devices( + host_id, + [_device(device_id, host_id)], + ) + for task_id in (first_task_id, second_task_id): + database.repository.enqueue_task( + ScheduledTask( + id=task_id, + goal="run task", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=assigned_at, + ) + ) + assert database.repository.assign_task( + task_id=first_task_id, + host_id=host_id, + device_id=device_id, + lease_id="lease-expired", + lease_expires_at=assigned_at + timedelta(minutes=1), + now=assigned_at, + ) + + assert device_id not in database.repository.list_reserved_device_ids( + now=after_expiry + ) + assert database.repository.assign_task( + task_id=second_task_id, + host_id=host_id, + device_id=device_id, + lease_id="lease-new", + lease_expires_at=after_expiry + timedelta(minutes=1), + now=after_expiry, + ) + finally: + database.close() diff --git a/tests/test_task_scheduler.py b/tests/test_task_scheduler.py index 3dbd9e5..c8c4a22 100644 --- a/tests/test_task_scheduler.py +++ b/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)