343 lines
11 KiB
Python
343 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from cloud.internal_api.models import (
|
|
AssignmentModel,
|
|
ClaimResponse,
|
|
DeviceEnrollmentResponse,
|
|
DeviceSnapshotModel,
|
|
HeartbeatResponse,
|
|
HostEnrollmentResponse,
|
|
HostTaskSubmissionResponse,
|
|
LeaseRenewalResponse,
|
|
TaskProgressModel,
|
|
TerminalResultResponse,
|
|
)
|
|
from host_agent.config import HostAgentConfig
|
|
from host_agent.progress import TaskProgressSnapshot
|
|
|
|
_VALID_STEP_STATUSES = frozenset({"running", "completed", "failed"})
|
|
|
|
|
|
class HostAgentAPIError(RuntimeError):
|
|
def __init__(self, status_code: int, detail: str) -> None:
|
|
super().__init__(
|
|
f"control plane request failed with status {status_code}: {detail}"
|
|
)
|
|
self.status_code = status_code
|
|
self.detail = detail
|
|
|
|
|
|
class StaleLeaseError(HostAgentAPIError):
|
|
pass
|
|
|
|
|
|
class HostTaskSubmissionUnknownError(RuntimeError):
|
|
"""Raised when a Host self-submission request's Cloud outcome is uncertain.
|
|
|
|
This is distinct from :class:`HostAgentAPIError` because the request may
|
|
have reached the control plane but the response was lost, the server
|
|
returned a 5xx, or the success payload was malformed. Retrying would
|
|
duplicate the create, so the caller must treat the task as unknown and
|
|
surface that to the operator.
|
|
"""
|
|
|
|
def __init__(self, reason: str) -> None:
|
|
super().__init__(f"host task submission outcome is unknown: {reason}")
|
|
self.reason = reason
|
|
|
|
|
|
class HostAgentEnrollmentClient:
|
|
def __init__(
|
|
self,
|
|
config: HostAgentConfig,
|
|
*,
|
|
http_client: httpx.Client | None = None,
|
|
sleep: Callable[[float], None] = time.sleep,
|
|
) -> None:
|
|
self.config = config
|
|
self._sleep = sleep
|
|
self._owns_client = http_client is None
|
|
self._client = http_client or httpx.Client(base_url=config.control_plane_url)
|
|
|
|
def enroll_host(
|
|
self,
|
|
*,
|
|
agent_instance_id: str,
|
|
host_token: str,
|
|
display_name: str | None,
|
|
) -> HostEnrollmentResponse:
|
|
response = self._request(
|
|
"POST",
|
|
"/internal/v1/enrollments",
|
|
token=None,
|
|
json={
|
|
"agent_instance_id": agent_instance_id,
|
|
"host_token": host_token,
|
|
"display_name": display_name,
|
|
},
|
|
)
|
|
return HostEnrollmentResponse.model_validate(response.json())
|
|
|
|
def enroll_device(
|
|
self,
|
|
*,
|
|
local_device_id: str,
|
|
driver_type: str,
|
|
name: str | None,
|
|
capability_tags: list[str],
|
|
) -> DeviceEnrollmentResponse:
|
|
if not self.config.host_id or not self.config.token:
|
|
raise HostAgentAPIError(0, "Host identity is unresolved")
|
|
response = self._request(
|
|
"POST",
|
|
f"/internal/v1/hosts/{self.config.host_id}/devices/enroll",
|
|
token=self.config.token,
|
|
json={
|
|
"local_device_id": local_device_id,
|
|
"driver_type": driver_type,
|
|
"name": name,
|
|
"capability_tags": list(capability_tags),
|
|
},
|
|
)
|
|
return DeviceEnrollmentResponse.model_validate(response.json())
|
|
|
|
def close(self) -> None:
|
|
if self._owns_client:
|
|
self._client.close()
|
|
|
|
def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
token: str | None,
|
|
json: dict[str, Any],
|
|
) -> httpx.Response:
|
|
headers = {"Authorization": f"Bearer {token}"} if token else {}
|
|
backoff = self.config.retry_backoff_seconds
|
|
for attempt in range(1, self.config.max_retry_attempts + 1):
|
|
try:
|
|
response = self._client.request(
|
|
method,
|
|
path,
|
|
json=json,
|
|
headers=headers,
|
|
)
|
|
except httpx.TransportError:
|
|
if attempt == self.config.max_retry_attempts:
|
|
raise
|
|
else:
|
|
if response.status_code < 500:
|
|
if response.is_success:
|
|
return response
|
|
_raise_api_error(response)
|
|
if attempt == self.config.max_retry_attempts:
|
|
_raise_api_error(response)
|
|
self._sleep(backoff)
|
|
backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
|
|
raise AssertionError("retry loop exited unexpectedly")
|
|
|
|
|
|
class HostAgentClient:
|
|
def __init__(
|
|
self,
|
|
config: HostAgentConfig,
|
|
*,
|
|
http_client: httpx.AsyncClient | None = None,
|
|
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
|
) -> None:
|
|
self.config = config
|
|
self._sleep = sleep
|
|
self._owns_client = http_client is None
|
|
self._client = http_client or httpx.AsyncClient(
|
|
base_url=config.control_plane_url,
|
|
headers={"Authorization": f"Bearer {config.token}"},
|
|
)
|
|
|
|
async def heartbeat(
|
|
self,
|
|
devices: list[DeviceSnapshotModel],
|
|
*,
|
|
address: str | None = None,
|
|
policy_revision: int = 0,
|
|
) -> HeartbeatResponse:
|
|
response = await self._request(
|
|
"PUT",
|
|
f"/internal/v1/hosts/{self.config.host_id}/heartbeat",
|
|
json={
|
|
"host_id": self.config.host_id,
|
|
"address": address,
|
|
"devices": [device.model_dump(mode="json") for device in devices],
|
|
"policy_revision": policy_revision,
|
|
"planner_transport": self.config.ai_planner_transport,
|
|
},
|
|
)
|
|
return HeartbeatResponse.model_validate(response.json())
|
|
|
|
async def submit_self_task(
|
|
self,
|
|
*,
|
|
goal: str,
|
|
device_id: str | None = None,
|
|
) -> HostTaskSubmissionResponse:
|
|
try:
|
|
response = await self._client.request(
|
|
"POST",
|
|
f"/internal/v1/hosts/{self.config.host_id}/tasks",
|
|
json={
|
|
"host_id": self.config.host_id,
|
|
"goal": goal,
|
|
"device_id": device_id,
|
|
},
|
|
headers={"Authorization": f"Bearer {self.config.token}"},
|
|
)
|
|
except httpx.TransportError as exc:
|
|
raise HostTaskSubmissionUnknownError(str(exc)) from exc
|
|
if response.status_code >= 500:
|
|
raise HostTaskSubmissionUnknownError(
|
|
f"control plane returned status {response.status_code}"
|
|
)
|
|
if not response.is_success:
|
|
_raise_api_error(response, stale_lease=False)
|
|
try:
|
|
payload = response.json()
|
|
except ValueError as exc:
|
|
raise HostTaskSubmissionUnknownError(
|
|
"control plane returned malformed success payload"
|
|
) from exc
|
|
try:
|
|
return HostTaskSubmissionResponse.model_validate(payload)
|
|
except Exception as exc:
|
|
raise HostTaskSubmissionUnknownError(
|
|
"control plane returned malformed success payload"
|
|
) from exc
|
|
|
|
async def claim(self) -> AssignmentModel | None:
|
|
response = await self._request(
|
|
"POST",
|
|
f"/internal/v1/hosts/{self.config.host_id}/assignments/claim",
|
|
json={
|
|
"host_id": self.config.host_id,
|
|
"timeout_seconds": self.config.poll_timeout_seconds,
|
|
},
|
|
timeout=self.config.poll_timeout_seconds + 5,
|
|
)
|
|
return ClaimResponse.model_validate(response.json()).assignment
|
|
|
|
async def renew(
|
|
self,
|
|
assignment: AssignmentModel,
|
|
*,
|
|
progress: TaskProgressSnapshot | None = None,
|
|
) -> LeaseRenewalResponse:
|
|
payload: dict[str, Any] = {
|
|
"host_id": self.config.host_id,
|
|
"task_id": assignment.task_id,
|
|
"attempt": assignment.attempt,
|
|
"lease_id": assignment.lease_id,
|
|
}
|
|
if progress is not None:
|
|
step_status = (
|
|
progress.step_status
|
|
if progress.step_status in _VALID_STEP_STATUSES
|
|
else "running"
|
|
)
|
|
payload["progress"] = TaskProgressModel(
|
|
step_index=max(progress.step_index, 0),
|
|
step_status=step_status,
|
|
summary=progress.summary[:500],
|
|
).model_dump(mode="json")
|
|
response = await self._request(
|
|
"POST",
|
|
(
|
|
f"/internal/v1/hosts/{self.config.host_id}/assignments/"
|
|
f"{assignment.task_id}/renew"
|
|
),
|
|
json=payload,
|
|
)
|
|
return LeaseRenewalResponse.model_validate(response.json())
|
|
|
|
async def report_result(
|
|
self,
|
|
assignment: AssignmentModel,
|
|
*,
|
|
status: str,
|
|
failure_reason: str | None = None,
|
|
result: dict[str, Any] | None = None,
|
|
) -> TerminalResultResponse:
|
|
response = await self._request(
|
|
"POST",
|
|
(
|
|
f"/internal/v1/hosts/{self.config.host_id}/assignments/"
|
|
f"{assignment.task_id}/result"
|
|
),
|
|
json={
|
|
"host_id": self.config.host_id,
|
|
"task_id": assignment.task_id,
|
|
"attempt": assignment.attempt,
|
|
"lease_id": assignment.lease_id,
|
|
"status": status,
|
|
"failure_reason": failure_reason,
|
|
"result": result,
|
|
},
|
|
)
|
|
return TerminalResultResponse.model_validate(response.json())
|
|
|
|
async def aclose(self) -> None:
|
|
if self._owns_client:
|
|
await self._client.aclose()
|
|
|
|
async def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
json: dict[str, Any],
|
|
timeout: float | None = None,
|
|
) -> httpx.Response:
|
|
backoff = self.config.retry_backoff_seconds
|
|
for attempt in range(1, self.config.max_retry_attempts + 1):
|
|
try:
|
|
response = await self._client.request(
|
|
method,
|
|
path,
|
|
json=json,
|
|
timeout=timeout,
|
|
headers={"Authorization": f"Bearer {self.config.token}"},
|
|
)
|
|
except httpx.TransportError:
|
|
if attempt == self.config.max_retry_attempts:
|
|
raise
|
|
else:
|
|
if response.status_code < 500:
|
|
if response.is_success:
|
|
return response
|
|
_raise_api_error(response, stale_lease=True)
|
|
if attempt == self.config.max_retry_attempts:
|
|
_raise_api_error(response, stale_lease=True)
|
|
|
|
await self._sleep(backoff)
|
|
backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
|
|
raise AssertionError("retry loop exited unexpectedly")
|
|
|
|
|
|
def _raise_api_error(response: httpx.Response, *, stale_lease: bool = False) -> None:
|
|
try:
|
|
payload = response.json()
|
|
except ValueError:
|
|
payload = {}
|
|
detail = payload.get("detail") or payload.get("code") or "request rejected"
|
|
error_type = (
|
|
StaleLeaseError
|
|
if stale_lease and response.status_code == 409
|
|
else HostAgentAPIError
|
|
)
|
|
raise error_type(response.status_code, str(detail))
|