100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from cloud.internal_api.models import AssignmentModel
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _CurrentAssignment:
|
|
task_id: str
|
|
device_id: str
|
|
goal: str | None
|
|
workflow_definition_id: str | None
|
|
started_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _LastHeartbeat:
|
|
ok: bool
|
|
device_count: int
|
|
at: datetime
|
|
|
|
|
|
class AgentStatusTracker:
|
|
def __init__(self, *, now: Callable[[], datetime] | None = None) -> None:
|
|
self._now = now or (lambda: datetime.now(UTC))
|
|
self._lock = threading.Lock()
|
|
self._current_assignment: _CurrentAssignment | None = None
|
|
self._last_heartbeat: _LastHeartbeat | None = None
|
|
self._host_policy: dict[str, Any] | None = None
|
|
|
|
def mark_assignment_started(self, assignment: AssignmentModel) -> None:
|
|
with self._lock:
|
|
self._current_assignment = _CurrentAssignment(
|
|
task_id=assignment.task_id,
|
|
device_id=assignment.device_id,
|
|
goal=assignment.goal,
|
|
workflow_definition_id=assignment.workflow_definition_id,
|
|
started_at=self._now(),
|
|
)
|
|
|
|
def mark_assignment_finished(self) -> None:
|
|
with self._lock:
|
|
self._current_assignment = None
|
|
|
|
def mark_heartbeat(self, *, ok: bool, device_count: int) -> None:
|
|
with self._lock:
|
|
self._last_heartbeat = _LastHeartbeat(
|
|
ok=ok,
|
|
device_count=device_count,
|
|
at=self._now(),
|
|
)
|
|
|
|
def mark_host_policy(self, policy: Any | None) -> None:
|
|
with self._lock:
|
|
self._host_policy = (
|
|
{
|
|
"revision": policy.revision,
|
|
"self_submission_enabled": policy.self_submission_enabled,
|
|
"max_active_tasks": policy.max_active_tasks,
|
|
"daily_token_budget": policy.daily_token_budget,
|
|
}
|
|
if policy is not None
|
|
else None
|
|
)
|
|
|
|
def snapshot(self) -> dict[str, Any]:
|
|
with self._lock:
|
|
current_assignment = self._current_assignment
|
|
last_heartbeat = self._last_heartbeat
|
|
return {
|
|
"current_assignment": (
|
|
{
|
|
"task_id": current_assignment.task_id,
|
|
"device_id": current_assignment.device_id,
|
|
"goal": current_assignment.goal,
|
|
"workflow_definition_id": current_assignment.workflow_definition_id,
|
|
"started_at": current_assignment.started_at.isoformat(),
|
|
}
|
|
if current_assignment is not None
|
|
else None
|
|
),
|
|
"last_heartbeat": (
|
|
{
|
|
"ok": last_heartbeat.ok,
|
|
"device_count": last_heartbeat.device_count,
|
|
"at": last_heartbeat.at.isoformat(),
|
|
}
|
|
if last_heartbeat is not None
|
|
else None
|
|
),
|
|
"host_policy": self._host_policy.copy() if self._host_policy else None,
|
|
}
|