q792602257
d3024b4810
- LeaseGuard gains an is_cancellation convenience property - ActiveAssignmentRunner marks the lease lost with a cancellation reason when a renewal response reports cancel_requested - AssignmentExecutor threads stop_reason through to TaskRunner/ WorkflowRunner and maps a cancellation-flavored stop to AssignmentExecutionResult.status = "cancelled" instead of "failed" - AssignmentProcessor forwards a three-way done/cancelled/failed status when reporting the terminal result - Add/extend tests across lease, assignment, processor, and client
122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Callable
|
|
from datetime import UTC, datetime
|
|
from threading import Event, Lock
|
|
from typing import Protocol
|
|
|
|
import httpx
|
|
|
|
from cloud.internal_api.models import AssignmentModel
|
|
from host_agent.assignment import AssignmentExecutionResult
|
|
from host_agent.client import HostAgentAPIError, HostAgentClient, StaleLeaseError
|
|
from host_agent.progress import TaskProgressSnapshot
|
|
from runtime.task import is_cancellation_reason
|
|
|
|
|
|
class InterruptibleAssignmentExecutor(Protocol):
|
|
def execute(
|
|
self,
|
|
assignment: AssignmentModel,
|
|
*,
|
|
should_stop: Callable[[], bool] | None = None,
|
|
stop_reason: Callable[[], str | None] | None = None,
|
|
) -> AssignmentExecutionResult: ...
|
|
|
|
def latest_progress(self) -> TaskProgressSnapshot | None: ...
|
|
|
|
|
|
class LeaseGuard:
|
|
def __init__(self) -> None:
|
|
self._lost = Event()
|
|
self._lock = Lock()
|
|
self._reason: str | None = None
|
|
|
|
@property
|
|
def reason(self) -> str | None:
|
|
with self._lock:
|
|
return self._reason
|
|
|
|
@property
|
|
def is_cancellation(self) -> bool:
|
|
return is_cancellation_reason(self.reason)
|
|
|
|
def is_lost(self) -> bool:
|
|
return self._lost.is_set()
|
|
|
|
def mark_lost(self, reason: str) -> None:
|
|
with self._lock:
|
|
if self._reason is None:
|
|
self._reason = reason
|
|
self._lost.set()
|
|
|
|
|
|
class ActiveAssignmentRunner:
|
|
def __init__(
|
|
self,
|
|
client: HostAgentClient,
|
|
executor: InterruptibleAssignmentExecutor,
|
|
*,
|
|
now: Callable[[], datetime] | None = None,
|
|
) -> None:
|
|
self.client = client
|
|
self.executor = executor
|
|
self._now = now or (lambda: datetime.now(UTC))
|
|
self._stop_requested = Event()
|
|
|
|
def request_stop(self) -> None:
|
|
self._stop_requested.set()
|
|
|
|
async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult:
|
|
guard = LeaseGuard()
|
|
execution = asyncio.create_task(
|
|
asyncio.to_thread(
|
|
self.executor.execute,
|
|
assignment,
|
|
should_stop=lambda: guard.is_lost() or self._stop_requested.is_set(),
|
|
stop_reason=lambda: guard.reason,
|
|
)
|
|
)
|
|
renewal = asyncio.create_task(
|
|
self._renew_while_running(assignment, execution, guard)
|
|
)
|
|
try:
|
|
return await execution
|
|
finally:
|
|
await renewal
|
|
|
|
async def _renew_while_running(
|
|
self,
|
|
assignment: AssignmentModel,
|
|
execution: asyncio.Task[AssignmentExecutionResult],
|
|
guard: LeaseGuard,
|
|
) -> None:
|
|
lease_expires_at = assignment.lease_expires_at
|
|
while not execution.done():
|
|
delay = max(
|
|
0.0,
|
|
(lease_expires_at - self._now()).total_seconds() / 3,
|
|
)
|
|
done, _ = await asyncio.wait({execution}, timeout=delay)
|
|
if done:
|
|
return
|
|
try:
|
|
response = await self.client.renew(
|
|
assignment, progress=self.executor.latest_progress()
|
|
)
|
|
except StaleLeaseError:
|
|
guard.mark_lost("lease rejected by control plane")
|
|
return
|
|
except HostAgentAPIError as exc:
|
|
guard.mark_lost(f"lease renewal rejected: {exc.detail}")
|
|
return
|
|
except httpx.TransportError:
|
|
guard.mark_lost("lease renewal failed after transport retries")
|
|
return
|
|
else:
|
|
if response.cancel_requested:
|
|
guard.mark_lost("cancellation requested by control plane")
|
|
return
|
|
lease_expires_at = response.lease_expires_at
|