feat(cloud-store): persist task lease history

This commit is contained in:
2026-07-12 17:15:09 +08:00
parent 52d2a86a4c
commit 0a9392b7ea
5 changed files with 155 additions and 7 deletions
@@ -17,7 +17,7 @@
## 3. Lease-Backed Scheduling ## 3. Lease-Backed Scheduling
- [ ] 3.1 Extend scheduled-task persistence with attempt count, lease id/expiry, terminal result, failure reason, and auditable attempt records. - [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. - [ ] 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.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.4 Implement lease renewal with host/task/attempt ownership validation and conflict responses for stale leases.
+2 -1
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING, Literal, Protocol from typing import TYPE_CHECKING, Any, Literal, Protocol
if TYPE_CHECKING: if TYPE_CHECKING:
from cloud.plugins import PluginManifest from cloud.plugins import PluginManifest
@@ -27,6 +27,7 @@ class TaskAttemptRecord:
created_at: datetime created_at: datetime
completed_at: datetime | None = None completed_at: datetime | None = None
failure_reason: str | None = None failure_reason: str | None = None
terminal_result: dict[str, Any] | None = None
@dataclass(frozen=True) @dataclass(frozen=True)
+10 -2
View File
@@ -10,7 +10,7 @@ from __future__ import annotations
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
from uuid import uuid4 from uuid import uuid4
from core.models import utc_now from core.models import utc_now
@@ -44,6 +44,12 @@ class ScheduledTask:
status: ScheduledTaskStatus = "queued" status: ScheduledTaskStatus = "queued"
assigned_device_id: str | None = None assigned_device_id: str | None = None
assigned_host_id: str | None = None assigned_host_id: str | None = None
attempt_count: int = 0
lease_id: str | None = None
lease_expires_at: datetime | None = None
terminal_result: dict[str, Any] | None = None
failure_reason: str | None = None
updated_at: datetime | None = None
created_at: datetime = field(default_factory=utc_now) created_at: datetime = field(default_factory=utc_now)
@@ -107,7 +113,9 @@ class TaskScheduler:
self.pool = pool self.pool = pool
self.store = store self.store = store
self.config = config self.config = config
self._strategies = dict(strategies) if strategies is not None else dict(DEFAULT_STRATEGIES) self._strategies = (
dict(strategies) if strategies is not None else dict(DEFAULT_STRATEGIES)
)
if config.default_assignment_strategy not in self._strategies: if config.default_assignment_strategy not in self._strategies:
raise UnknownAssignmentStrategyError( raise UnknownAssignmentStrategyError(
f"unknown assignment strategy {config.default_assignment_strategy!r}; " f"unknown assignment strategy {config.default_assignment_strategy!r}; "
@@ -6,7 +6,7 @@ from datetime import datetime
from typing import Any from typing import Any
from sqlalchemy import Engine, delete, func, select from sqlalchemy import Engine, delete, func, select
from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.orm import sessionmaker
from cloud.db_models import ( from cloud.db_models import (
Base, Base,
@@ -14,6 +14,7 @@ from cloud.db_models import (
PluginRow, PluginRow,
PooledDeviceRow, PooledDeviceRow,
ScheduledTaskRow, ScheduledTaskRow,
TaskAttemptRow,
) )
from core.models import utc_now from core.models import utc_now
@@ -119,6 +120,18 @@ class SQLAlchemyCloudRepository:
status=task.status, status=task.status,
assigned_device_id=task.assigned_device_id, assigned_device_id=task.assigned_device_id,
assigned_host_id=task.assigned_host_id, assigned_host_id=task.assigned_host_id,
attempt_count=task.attempt_count,
lease_id=task.lease_id,
lease_expires_at=(
_iso(task.lease_expires_at) if task.lease_expires_at else None
),
failure_reason=task.failure_reason,
result_json=(
json.dumps(task.terminal_result, ensure_ascii=False)
if task.terminal_result is not None
else None
),
updated_at=_iso(task.updated_at) if task.updated_at else None,
created_at=_iso(task.created_at), created_at=_iso(task.created_at),
) )
) )
@@ -194,6 +207,15 @@ class SQLAlchemyCloudRepository:
row = session.get(PluginRow, name) row = session.get(PluginRow, name)
return _plugin_from_row(row) if row else None return _plugin_from_row(row) if row else None
def list_task_attempts(self, task_id: str) -> list[Any]:
with self._sessions() as session:
rows = session.scalars(
select(TaskAttemptRow)
.where(TaskAttemptRow.task_id == task_id)
.order_by(TaskAttemptRow.attempt)
).all()
return [_task_attempt_from_row(row) for row in rows]
def health_check(self) -> None: def health_check(self) -> None:
with self._sessions() as session: with self._sessions() as session:
session.execute(select(1)) session.execute(select(1))
@@ -230,7 +252,7 @@ def _device_from_row(row: PooledDeviceRow) -> Any:
try: try:
tags = list(json.loads(row.capability_tags_json)) tags = list(json.loads(row.capability_tags_json))
except (TypeError, ValueError): except TypeError, ValueError:
tags = [] tags = []
return PooledDevice( return PooledDevice(
device_id=row.device_id, device_id=row.device_id,
@@ -247,8 +269,9 @@ def _task_from_row(row: ScheduledTaskRow) -> Any:
try: try:
constraints_data = json.loads(row.constraints_json) constraints_data = json.loads(row.constraints_json)
except (TypeError, ValueError): except TypeError, ValueError:
constraints_data = {} constraints_data = {}
terminal_result = _parse_json_object(row.result_json)
return ScheduledTask( return ScheduledTask(
id=row.id, id=row.id,
goal=row.goal, goal=row.goal,
@@ -260,10 +283,44 @@ def _task_from_row(row: ScheduledTaskRow) -> Any:
status=row.status, status=row.status,
assigned_device_id=row.assigned_device_id, assigned_device_id=row.assigned_device_id,
assigned_host_id=row.assigned_host_id, assigned_host_id=row.assigned_host_id,
attempt_count=row.attempt_count,
lease_id=row.lease_id,
lease_expires_at=_parse_dt(row.lease_expires_at),
terminal_result=terminal_result,
failure_reason=row.failure_reason,
updated_at=_parse_dt(row.updated_at),
created_at=_parse_dt(row.created_at) or utc_now(), created_at=_parse_dt(row.created_at) or utc_now(),
) )
def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
from cloud.repository import TaskAttemptRecord
return TaskAttemptRecord(
task_id=row.task_id,
attempt=row.attempt,
lease_id=row.lease_id,
host_id=row.host_id,
device_id=row.device_id,
status=row.status,
lease_expires_at=_parse_dt(row.lease_expires_at) or utc_now(),
created_at=_parse_dt(row.created_at) or utc_now(),
completed_at=_parse_dt(row.completed_at),
failure_reason=row.failure_reason,
terminal_result=_parse_json_object(row.result_json),
)
def _parse_json_object(value: str | None) -> dict[str, Any] | None:
if value is None:
return None
try:
parsed = json.loads(value)
except TypeError, ValueError:
return None
return parsed if isinstance(parsed, dict) else None
def _plugin_from_row(row: PluginRow) -> tuple[Any, bool]: def _plugin_from_row(row: PluginRow) -> tuple[Any, bool]:
from cloud.plugins import PluginManifest from cloud.plugins import PluginManifest
+82
View File
@@ -8,8 +8,10 @@ from uuid import uuid4
import pytest import pytest
from sqlalchemy import event from sqlalchemy import event
from sqlalchemy.orm import Session
from cloud.database import CloudDatabase from cloud.database import CloudDatabase
from cloud.db_models import TaskAttemptRow
from cloud.plugins import PluginManifest from cloud.plugins import PluginManifest
from cloud.pool import PooledDevice from cloud.pool import PooledDevice
from cloud.repository import CloudRepository, LeasedAssignment, TaskAttemptRecord from cloud.repository import CloudRepository, LeasedAssignment, TaskAttemptRecord
@@ -196,3 +198,83 @@ def test_repository_state_survives_application_restart(database_url: str) -> Non
assert device == _device(device_id, host_id) assert device == _device(device_id, host_id)
finally: finally:
restarted_process.close() restarted_process.close()
def test_task_lease_and_terminal_fields_round_trip(database_url: str) -> None:
database = CloudDatabase(database_url)
task_id = _unique_id("lease-task")
created_at = datetime(2026, 7, 12, 1, 0, tzinfo=UTC)
lease_expires_at = datetime(2026, 7, 12, 1, 5, tzinfo=UTC)
updated_at = datetime(2026, 7, 12, 1, 1, tzinfo=UTC)
task = ScheduledTask(
id=task_id,
goal="capture diagnostics",
workflow_definition_id=None,
constraints=TaskConstraints(),
status="failed",
assigned_device_id="device-a",
assigned_host_id="host-a",
attempt_count=2,
lease_id="lease-a",
lease_expires_at=lease_expires_at,
terminal_result={"steps": 3, "status": "failed"},
failure_reason="device disconnected",
updated_at=updated_at,
created_at=created_at,
)
try:
database.repository.enqueue_task(task)
assert database.repository.get_task(task_id) == task
finally:
database.close()
def test_task_attempt_history_is_ordered_and_complete(database_url: str) -> None:
database = CloudDatabase(database_url)
task_id = _unique_id("attempt-task")
created_at = datetime(2026, 7, 12, 1, 0, tzinfo=UTC)
lease_expires_at = datetime(2026, 7, 12, 1, 5, tzinfo=UTC)
try:
with Session(database.engine) as session, session.begin():
session.add_all(
[
TaskAttemptRow(
task_id=task_id,
attempt=2,
lease_id="lease-2",
host_id="host-b",
device_id="device-b",
status="failed",
lease_expires_at=lease_expires_at.isoformat(),
created_at=created_at.isoformat(),
completed_at=lease_expires_at.isoformat(),
failure_reason="execution failed",
result_json='{"exit_code": 1}',
),
TaskAttemptRow(
task_id=task_id,
attempt=1,
lease_id="lease-1",
host_id="host-a",
device_id="device-a",
status="expired",
lease_expires_at=lease_expires_at.isoformat(),
created_at=created_at.isoformat(),
completed_at=lease_expires_at.isoformat(),
failure_reason="lease expired",
result_json=None,
),
]
)
attempts = database.repository.list_task_attempts(task_id)
assert [attempt.attempt for attempt in attempts] == [1, 2]
assert attempts[0].status == "expired"
assert attempts[1].terminal_result == {"exit_code": 1}
assert attempts[1].failure_reason == "execution failed"
finally:
database.close()