feat(host-agent): stop assignment execution collaboratively on cancellation
- 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
This commit is contained in:
@@ -12,6 +12,7 @@ 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):
|
||||
@@ -20,6 +21,7 @@ class InterruptibleAssignmentExecutor(Protocol):
|
||||
assignment: AssignmentModel,
|
||||
*,
|
||||
should_stop: Callable[[], bool] | None = None,
|
||||
stop_reason: Callable[[], str | None] | None = None,
|
||||
) -> AssignmentExecutionResult: ...
|
||||
|
||||
def latest_progress(self) -> TaskProgressSnapshot | None: ...
|
||||
@@ -36,6 +38,10 @@ class LeaseGuard:
|
||||
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()
|
||||
|
||||
@@ -69,6 +75,7 @@ class ActiveAssignmentRunner:
|
||||
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(
|
||||
@@ -108,4 +115,7 @@ class ActiveAssignmentRunner:
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user