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
74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Literal, Protocol
|
|
|
|
from cloud.internal_api.models import AssignmentModel
|
|
from host_agent.assignment import AssignmentExecutionResult
|
|
from host_agent.client import HostAgentClient
|
|
from host_agent.status import AgentStatusTracker
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
|
|
class ActiveAssignmentExecutor(Protocol):
|
|
async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult: ...
|
|
|
|
def request_stop(self) -> None: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssignmentProcessingResult:
|
|
execution: AssignmentExecutionResult
|
|
report_status: Literal["recorded", "already_recorded"]
|
|
|
|
|
|
class AssignmentProcessor:
|
|
def __init__(
|
|
self,
|
|
client: HostAgentClient,
|
|
active_executor: ActiveAssignmentExecutor,
|
|
*,
|
|
status_tracker: AgentStatusTracker | None = None,
|
|
on_result: Callable[[AssignmentModel, AssignmentProcessingResult], None]
|
|
| None = None,
|
|
) -> None:
|
|
self.client = client
|
|
self.active_executor = active_executor
|
|
self.status_tracker = status_tracker
|
|
self.on_result = on_result
|
|
|
|
def request_stop(self) -> None:
|
|
self.active_executor.request_stop()
|
|
|
|
async def process(self, assignment: AssignmentModel) -> AssignmentProcessingResult:
|
|
if self.status_tracker is not None:
|
|
self.status_tracker.mark_assignment_started(assignment)
|
|
try:
|
|
execution = await self.active_executor.run(assignment)
|
|
if execution.status in {"done", "cancelled"}:
|
|
status = execution.status
|
|
else:
|
|
status = "failed"
|
|
failure_reason = execution.failure_reason if status != "done" else None
|
|
response = await self.client.report_result(
|
|
assignment,
|
|
status=status,
|
|
failure_reason=failure_reason,
|
|
result=dict(execution.metadata),
|
|
)
|
|
result = AssignmentProcessingResult(
|
|
execution=execution,
|
|
report_status=response.status,
|
|
)
|
|
finally:
|
|
if self.status_tracker is not None:
|
|
self.status_tracker.mark_assignment_finished()
|
|
if self.on_result is not None:
|
|
try:
|
|
self.on_result(assignment, result)
|
|
except Exception:
|
|
pass
|
|
return result
|