|
|
@@ -0,0 +1,796 @@
|
|
|
+"""Durable background-operation lifecycle control.
|
|
|
+
|
|
|
+The controller owns process-local scheduling and operation-level state only.
|
|
|
+Task/Attempt/Validation execution remains in :mod:`coordinator` and is reached
|
|
|
+through the injected ``execute_claimed`` callback.
|
|
|
+"""
|
|
|
+
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import asyncio
|
|
|
+from dataclasses import replace
|
|
|
+from datetime import datetime, timezone
|
|
|
+from functools import partial
|
|
|
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Protocol, Sequence
|
|
|
+
|
|
|
+from .config import OrchestrationConfig
|
|
|
+from .errors import TaskConflict
|
|
|
+from .models import (
|
|
|
+ AttemptStatus,
|
|
|
+ BackgroundOperation,
|
|
|
+ ExecutionStats,
|
|
|
+ FailureCode,
|
|
|
+ OperationKind,
|
|
|
+ OperationStatus,
|
|
|
+ TaskCycleResult,
|
|
|
+ TaskLedger,
|
|
|
+ TaskStatus,
|
|
|
+ TaskAttempt,
|
|
|
+ ValidationReport,
|
|
|
+ ValidationRunStatus,
|
|
|
+ new_id,
|
|
|
+ utc_now,
|
|
|
+)
|
|
|
+from .protocols import AgentExecutor, TaskStore
|
|
|
+from .runtime import OperationRuntime
|
|
|
+
|
|
|
+
|
|
|
+class Mutate(Protocol):
|
|
|
+ async def __call__(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ event_type: str,
|
|
|
+ mutator: Callable[[TaskLedger], Dict[str, Any]],
|
|
|
+ idempotency_key: Optional[str] = None,
|
|
|
+ *,
|
|
|
+ idempotency_payload: Optional[Any] = None,
|
|
|
+ operation_id: Optional[str] = None,
|
|
|
+ ) -> Dict[str, Any]: ...
|
|
|
+
|
|
|
+
|
|
|
+ExecuteClaimed = Callable[
|
|
|
+ [BackgroundOperation],
|
|
|
+ Awaitable[Sequence[TaskCycleResult]],
|
|
|
+]
|
|
|
+Fingerprint = Callable[[str, Any], str]
|
|
|
+
|
|
|
+
|
|
|
+class OperationController:
|
|
|
+ """Own BackgroundOperation persistence, recovery, stop and scheduling."""
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ task_store: TaskStore,
|
|
|
+ mutate: Mutate,
|
|
|
+ execute_claimed: ExecuteClaimed,
|
|
|
+ get_config: Callable[[], OrchestrationConfig],
|
|
|
+ get_executor: Callable[[], Optional[AgentExecutor]],
|
|
|
+ fingerprint: Fingerprint,
|
|
|
+ ) -> None:
|
|
|
+ self.task_store = task_store
|
|
|
+ self._mutate = mutate
|
|
|
+ self._execute_claimed = execute_claimed
|
|
|
+ self._get_config = get_config
|
|
|
+ self._get_executor = get_executor
|
|
|
+ self._fingerprint = fingerprint
|
|
|
+ self.runtime = OperationRuntime()
|
|
|
+
|
|
|
+ async def start(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ kind: OperationKind | str,
|
|
|
+ *,
|
|
|
+ task_ids: Optional[Sequence[str]] = None,
|
|
|
+ worker_presets: Optional[Sequence[str]] = None,
|
|
|
+ task_id: Optional[str] = None,
|
|
|
+ attempt_id: Optional[str] = None,
|
|
|
+ idempotency_key: Optional[str] = None,
|
|
|
+ deadline_at: Optional[str] = None,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ if not self._get_executor():
|
|
|
+ raise RuntimeError("TaskCoordinator has no AgentExecutor")
|
|
|
+ operation_kind, requested_task_ids, presets, request = _normalize_start_request(
|
|
|
+ kind,
|
|
|
+ task_ids,
|
|
|
+ task_id,
|
|
|
+ worker_presets,
|
|
|
+ attempt_id,
|
|
|
+ deadline_at,
|
|
|
+ self._get_config().default_worker_preset,
|
|
|
+ )
|
|
|
+ operation_id = new_id()
|
|
|
+ fingerprint = self._fingerprint("start_operation", request)
|
|
|
+ command_key = (
|
|
|
+ f"operation.{operation_kind.value}:{idempotency_key}"
|
|
|
+ if idempotency_key
|
|
|
+ else None
|
|
|
+ )
|
|
|
+ try:
|
|
|
+ result = await self._mutate(
|
|
|
+ root_trace_id,
|
|
|
+ "operation_started",
|
|
|
+ partial(
|
|
|
+ _create_operation,
|
|
|
+ operation_id=operation_id,
|
|
|
+ root_trace_id=root_trace_id,
|
|
|
+ operation_kind=operation_kind,
|
|
|
+ request=request,
|
|
|
+ fingerprint=fingerprint,
|
|
|
+ task_ids=requested_task_ids,
|
|
|
+ ),
|
|
|
+ command_key,
|
|
|
+ idempotency_payload=request,
|
|
|
+ operation_id=operation_id,
|
|
|
+ )
|
|
|
+ except TaskConflict:
|
|
|
+ if not command_key:
|
|
|
+ raise
|
|
|
+ await self._check_start_conflict(
|
|
|
+ root_trace_id,
|
|
|
+ command_key,
|
|
|
+ requested_task_ids,
|
|
|
+ presets,
|
|
|
+ )
|
|
|
+ raise
|
|
|
+
|
|
|
+ operation = await self.get(root_trace_id, result["operation_id"])
|
|
|
+ self.schedule(operation)
|
|
|
+ return operation
|
|
|
+
|
|
|
+ async def _check_start_conflict(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ command_key: str,
|
|
|
+ task_ids: Sequence[str],
|
|
|
+ presets: Sequence[str],
|
|
|
+ ) -> None:
|
|
|
+ ledger = await self.task_store.load(root_trace_id)
|
|
|
+ command = ledger.command_records.get(f"{root_trace_id}:{command_key}")
|
|
|
+ prior = ledger.operations.get(command.operation_id or "") if command else None
|
|
|
+ if not prior and command:
|
|
|
+ prior_id = command.result_ref.get("operation_id")
|
|
|
+ prior = ledger.operations.get(str(prior_id)) if prior_id else None
|
|
|
+ if prior and prior.task_ids != list(task_ids):
|
|
|
+ raise TaskConflict(
|
|
|
+ "Idempotency key is bound to different task_ids"
|
|
|
+ ) from None
|
|
|
+ if prior and list(prior.request.get("worker_presets", [])) != list(presets):
|
|
|
+ raise TaskConflict(
|
|
|
+ "Idempotency key is bound to different worker_presets"
|
|
|
+ ) from None
|
|
|
+
|
|
|
+ async def get(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_id: str,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ ledger = await self.task_store.load(root_trace_id)
|
|
|
+ try:
|
|
|
+ return ledger.operations[operation_id]
|
|
|
+ except KeyError as exc:
|
|
|
+ raise ValueError(f"Operation not found: {operation_id}") from exc
|
|
|
+
|
|
|
+ async def wait(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_id: str,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ operation = await self.get(root_trace_id, operation_id)
|
|
|
+ if operation.status == OperationStatus.PENDING:
|
|
|
+ self.schedule(operation)
|
|
|
+ await self.runtime.wait(root_trace_id, operation_id)
|
|
|
+ while True:
|
|
|
+ operation = await self.get(root_trace_id, operation_id)
|
|
|
+ if operation.status in {
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ OperationStatus.STOPPED,
|
|
|
+ }:
|
|
|
+ return operation
|
|
|
+ await asyncio.sleep(0.05)
|
|
|
+
|
|
|
+ async def stop(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_id: str,
|
|
|
+ idempotency_key: Optional[str] = None,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ """Persist stop intent before signalling any process-local Agent task."""
|
|
|
+
|
|
|
+ current = await self.get(root_trace_id, operation_id)
|
|
|
+ if current.status in {
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ OperationStatus.STOPPED,
|
|
|
+ }:
|
|
|
+ return current
|
|
|
+
|
|
|
+ def request_stop(ledger: TaskLedger) -> Dict[str, Any]:
|
|
|
+ operation = ledger.operations[operation_id]
|
|
|
+ if operation.status in {
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ OperationStatus.STOPPED,
|
|
|
+ }:
|
|
|
+ return {"operation_id": operation_id, "trace_ids": []}
|
|
|
+ operation.status = OperationStatus.STOP_REQUESTED
|
|
|
+ operation.execution_epoch += 1
|
|
|
+ operation.updated_at = utc_now()
|
|
|
+ trace_ids = [
|
|
|
+ ledger.attempts[item].worker_trace_id
|
|
|
+ for item in operation.attempt_ids
|
|
|
+ if item in ledger.attempts
|
|
|
+ and ledger.attempts[item].status == AttemptStatus.RUNNING
|
|
|
+ ]
|
|
|
+ trace_ids.extend(
|
|
|
+ ledger.validations[item].validator_trace_id
|
|
|
+ for item in operation.validation_ids
|
|
|
+ if item in ledger.validations
|
|
|
+ and ledger.validations[item].status == ValidationRunStatus.RUNNING
|
|
|
+ )
|
|
|
+ return {"operation_id": operation_id, "trace_ids": trace_ids}
|
|
|
+
|
|
|
+ result = await self._mutate(
|
|
|
+ root_trace_id,
|
|
|
+ "operation_stop_requested",
|
|
|
+ request_stop,
|
|
|
+ f"operation.stop:{idempotency_key}" if idempotency_key else None,
|
|
|
+ idempotency_payload={"operation_id": operation_id},
|
|
|
+ operation_id=operation_id,
|
|
|
+ )
|
|
|
+ executor = self._get_executor()
|
|
|
+ stop = getattr(executor, "stop", None)
|
|
|
+ signal_task: Optional[asyncio.Future[Any]] = None
|
|
|
+ if stop:
|
|
|
+ signal_task = asyncio.ensure_future(
|
|
|
+ asyncio.gather(
|
|
|
+ *(stop(trace_id) for trace_id in result["trace_ids"]),
|
|
|
+ return_exceptions=True,
|
|
|
+ )
|
|
|
+ )
|
|
|
+ grace = self._get_config().stop_grace_seconds
|
|
|
+ cancel_task = asyncio.create_task(
|
|
|
+ self.runtime.cancel(root_trace_id, operation_id, grace)
|
|
|
+ )
|
|
|
+ if signal_task:
|
|
|
+ await asyncio.sleep(0)
|
|
|
+ try:
|
|
|
+ await asyncio.wait_for(asyncio.shield(signal_task), timeout=grace)
|
|
|
+ except asyncio.TimeoutError:
|
|
|
+ signal_task.cancel()
|
|
|
+ await asyncio.gather(signal_task, return_exceptions=True)
|
|
|
+ await cancel_task
|
|
|
+ return await self._normalize_stopped(root_trace_id, operation_id)
|
|
|
+
|
|
|
+ async def _normalize_stopped(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_id: str,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ current = await self.get(root_trace_id, operation_id)
|
|
|
+ if current.status in {
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ OperationStatus.STOPPED,
|
|
|
+ }:
|
|
|
+ return current
|
|
|
+
|
|
|
+ def normalize(ledger: TaskLedger) -> Dict[str, Any]:
|
|
|
+ operation = ledger.operations[operation_id]
|
|
|
+ if operation.status in {OperationStatus.COMPLETED, OperationStatus.FAILED}:
|
|
|
+ return {"operation_id": operation_id, "status": operation.status.value}
|
|
|
+ for attempt_id in operation.attempt_ids:
|
|
|
+ attempt = ledger.attempts.get(attempt_id)
|
|
|
+ if attempt and attempt.status == AttemptStatus.RUNNING and attempt.started_at:
|
|
|
+ attempt.status = AttemptStatus.STOPPED
|
|
|
+ attempt.completed_at = utc_now()
|
|
|
+ attempt.error = attempt.error or "Operation stopped"
|
|
|
+ attempt.execution_stats = _failure_stats(
|
|
|
+ attempt.execution_stats, FailureCode.STOPPED
|
|
|
+ )
|
|
|
+ task = ledger.tasks[attempt.task_id]
|
|
|
+ if task.status == TaskStatus.RUNNING:
|
|
|
+ task.status = TaskStatus.NEEDS_REPLAN
|
|
|
+ task.updated_at = utc_now()
|
|
|
+ for validation_id in operation.validation_ids:
|
|
|
+ validation = ledger.validations.get(validation_id)
|
|
|
+ if (
|
|
|
+ validation
|
|
|
+ and validation.status == ValidationRunStatus.RUNNING
|
|
|
+ and validation.started_at
|
|
|
+ ):
|
|
|
+ validation.status = ValidationRunStatus.STOPPED
|
|
|
+ validation.completed_at = utc_now()
|
|
|
+ validation.error = validation.error or "Operation stopped"
|
|
|
+ validation.execution_stats = _failure_stats(
|
|
|
+ validation.execution_stats, FailureCode.STOPPED
|
|
|
+ )
|
|
|
+ task = ledger.tasks[validation.task_id]
|
|
|
+ if task.status == TaskStatus.VALIDATING:
|
|
|
+ task.status = TaskStatus.NEEDS_REPLAN
|
|
|
+ task.updated_at = utc_now()
|
|
|
+ operation.status = OperationStatus.STOPPED
|
|
|
+ operation.completed_at = utc_now()
|
|
|
+ operation.updated_at = operation.completed_at
|
|
|
+ return {"operation_id": operation_id, "status": operation.status.value}
|
|
|
+
|
|
|
+ await self._mutate(
|
|
|
+ root_trace_id,
|
|
|
+ "operation_stopped",
|
|
|
+ normalize,
|
|
|
+ operation_id=operation_id,
|
|
|
+ )
|
|
|
+ return await self.get(root_trace_id, operation_id)
|
|
|
+
|
|
|
+ async def resume(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_id: str,
|
|
|
+ idempotency_key: Optional[str] = None,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ """Resume only stages that never started, or validation after submit."""
|
|
|
+
|
|
|
+ def resume_state(ledger: TaskLedger) -> Dict[str, Any]:
|
|
|
+ operation = ledger.operations[operation_id]
|
|
|
+ if operation.status == OperationStatus.PENDING:
|
|
|
+ return {"operation_id": operation_id}
|
|
|
+ if operation.status != OperationStatus.STOPPED:
|
|
|
+ raise TaskConflict("Only a stopped operation can be resumed")
|
|
|
+ unsafe: List[str] = []
|
|
|
+ for attempt_id in operation.attempt_ids:
|
|
|
+ attempt = ledger.attempts[attempt_id]
|
|
|
+ if attempt.status in {
|
|
|
+ AttemptStatus.FAILED,
|
|
|
+ AttemptStatus.EXPIRED,
|
|
|
+ } or (attempt.status == AttemptStatus.STOPPED and attempt.started_at):
|
|
|
+ unsafe.append(attempt_id)
|
|
|
+ for validation_id in operation.validation_ids:
|
|
|
+ validation = ledger.validations[validation_id]
|
|
|
+ if validation.status in {
|
|
|
+ ValidationRunStatus.ERROR,
|
|
|
+ ValidationRunStatus.EXPIRED,
|
|
|
+ } or (
|
|
|
+ validation.status == ValidationRunStatus.STOPPED
|
|
|
+ and validation.started_at
|
|
|
+ ):
|
|
|
+ unsafe.append(validation_id)
|
|
|
+ if unsafe:
|
|
|
+ raise TaskConflict(f"resume_not_safe: stages already started: {unsafe}")
|
|
|
+ operation.status = OperationStatus.PENDING
|
|
|
+ operation.completed_at = None
|
|
|
+ operation.error = None
|
|
|
+ operation.updated_at = utc_now()
|
|
|
+ return {"operation_id": operation_id}
|
|
|
+
|
|
|
+ await self._mutate(
|
|
|
+ root_trace_id,
|
|
|
+ "operation_resumed",
|
|
|
+ resume_state,
|
|
|
+ f"operation.resume:{idempotency_key}" if idempotency_key else None,
|
|
|
+ idempotency_payload={"operation_id": operation_id},
|
|
|
+ operation_id=operation_id,
|
|
|
+ )
|
|
|
+ operation = await self.get(root_trace_id, operation_id)
|
|
|
+ self.schedule(operation)
|
|
|
+ return operation
|
|
|
+
|
|
|
+ async def recover(self, root_trace_id: str) -> List[BackgroundOperation]:
|
|
|
+ """Reconcile crash state without automatically invoking any model."""
|
|
|
+
|
|
|
+ recoverable = await self.task_store.list_recoverable(root_trace_id)
|
|
|
+ recovered: List[BackgroundOperation] = []
|
|
|
+ for candidate in recoverable:
|
|
|
+ await self._mutate(
|
|
|
+ root_trace_id,
|
|
|
+ "operation_recovered",
|
|
|
+ partial(_recover_operation_state, operation_id=candidate.operation_id),
|
|
|
+ operation_id=candidate.operation_id,
|
|
|
+ )
|
|
|
+ recovered.append(await self.get(root_trace_id, candidate.operation_id))
|
|
|
+ return recovered
|
|
|
+
|
|
|
+ def schedule(self, operation: BackgroundOperation) -> None:
|
|
|
+ self.runtime.schedule(operation, lambda: self._run(operation))
|
|
|
+
|
|
|
+ async def _run(self, operation: BackgroundOperation) -> None:
|
|
|
+ try:
|
|
|
+ await self.advance(operation.root_trace_id, operation.operation_id)
|
|
|
+ except asyncio.CancelledError:
|
|
|
+ raise
|
|
|
+ except Exception as exc:
|
|
|
+ await self.finish(
|
|
|
+ operation.root_trace_id,
|
|
|
+ operation.operation_id,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ error=f"Operation failed: {type(exc).__name__}: {exc}",
|
|
|
+ execution_epoch=operation.execution_epoch,
|
|
|
+ )
|
|
|
+
|
|
|
+ async def advance(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_id: str,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ """Claim and advance an operation through its durable stage state."""
|
|
|
+
|
|
|
+ def claim(ledger: TaskLedger) -> Dict[str, Any]:
|
|
|
+ operation = ledger.operations[operation_id]
|
|
|
+ if operation.status in {
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ OperationStatus.STOPPED,
|
|
|
+ }:
|
|
|
+ return {"operation_id": operation_id, "execute": False}
|
|
|
+ if operation.status == OperationStatus.RUNNING:
|
|
|
+ return {"operation_id": operation_id, "execute": False}
|
|
|
+ if operation.status == OperationStatus.STOP_REQUESTED:
|
|
|
+ return {
|
|
|
+ "operation_id": operation_id,
|
|
|
+ "execute": False,
|
|
|
+ "normalize_stop": True,
|
|
|
+ }
|
|
|
+ if deadline_expired(operation.deadline_at):
|
|
|
+ operation.status = OperationStatus.FAILED
|
|
|
+ operation.error = "Operation deadline exceeded before execution"
|
|
|
+ operation.completed_at = utc_now()
|
|
|
+ operation.updated_at = operation.completed_at
|
|
|
+ return {"operation_id": operation_id, "execute": False}
|
|
|
+ operation.status = OperationStatus.RUNNING
|
|
|
+ operation.execution_epoch += 1
|
|
|
+ operation.started_at = operation.started_at or utc_now()
|
|
|
+ operation.updated_at = utc_now()
|
|
|
+ return {
|
|
|
+ "operation_id": operation_id,
|
|
|
+ "execute": True,
|
|
|
+ "execution_epoch": operation.execution_epoch,
|
|
|
+ }
|
|
|
+
|
|
|
+ claimed = await self._mutate(
|
|
|
+ root_trace_id,
|
|
|
+ "operation_claimed",
|
|
|
+ claim,
|
|
|
+ operation_id=operation_id,
|
|
|
+ )
|
|
|
+ if not claimed["execute"]:
|
|
|
+ if claimed.get("normalize_stop"):
|
|
|
+ return await self._normalize_stopped(root_trace_id, operation_id)
|
|
|
+ return await self.get(root_trace_id, operation_id)
|
|
|
+ claimed_epoch = int(claimed["execution_epoch"])
|
|
|
+ try:
|
|
|
+ operation = await self.get(root_trace_id, operation_id)
|
|
|
+ results = list(await self._execute_claimed(operation))
|
|
|
+ return await self.finish(
|
|
|
+ root_trace_id,
|
|
|
+ operation_id,
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ results=results,
|
|
|
+ execution_epoch=claimed_epoch,
|
|
|
+ )
|
|
|
+ except asyncio.CancelledError:
|
|
|
+ raise
|
|
|
+ except Exception as exc:
|
|
|
+ return await self.finish(
|
|
|
+ root_trace_id,
|
|
|
+ operation_id,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ error=f"Operation failed: {type(exc).__name__}: {exc}",
|
|
|
+ execution_epoch=claimed_epoch,
|
|
|
+ )
|
|
|
+
|
|
|
+ async def finish(
|
|
|
+ self,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_id: str,
|
|
|
+ status: OperationStatus,
|
|
|
+ *,
|
|
|
+ results: Optional[Sequence[TaskCycleResult]] = None,
|
|
|
+ error: Optional[str] = None,
|
|
|
+ execution_epoch: Optional[int] = None,
|
|
|
+ ) -> BackgroundOperation:
|
|
|
+ result_list = list(results or [])
|
|
|
+
|
|
|
+ def finish_state(ledger: TaskLedger) -> Dict[str, Any]:
|
|
|
+ operation = ledger.operations[operation_id]
|
|
|
+ if operation.status in {
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ OperationStatus.STOPPED,
|
|
|
+ }:
|
|
|
+ return {
|
|
|
+ "operation_id": operation_id,
|
|
|
+ "status": operation.status.value,
|
|
|
+ "result_ref": operation.result_ref,
|
|
|
+ "error": operation.error,
|
|
|
+ }
|
|
|
+ if (
|
|
|
+ execution_epoch is not None
|
|
|
+ and operation.execution_epoch != execution_epoch
|
|
|
+ ):
|
|
|
+ return {
|
|
|
+ "operation_id": operation_id,
|
|
|
+ "status": operation.status.value,
|
|
|
+ "result_ref": operation.result_ref,
|
|
|
+ "error": operation.error,
|
|
|
+ }
|
|
|
+ final_status = (
|
|
|
+ OperationStatus.STOPPED
|
|
|
+ if operation.status == OperationStatus.STOP_REQUESTED
|
|
|
+ else status
|
|
|
+ )
|
|
|
+ operation.status = final_status
|
|
|
+ operation.error = error
|
|
|
+ for item in result_list:
|
|
|
+ if item.attempt_id and item.attempt_id not in operation.attempt_ids:
|
|
|
+ operation.attempt_ids.append(item.attempt_id)
|
|
|
+ if item.validation_id and item.validation_id not in operation.validation_ids:
|
|
|
+ operation.validation_ids.append(item.validation_id)
|
|
|
+ operation.result_ref = {
|
|
|
+ "task_ids": [item.task_id for item in result_list],
|
|
|
+ "attempt_ids": [item.attempt_id for item in result_list],
|
|
|
+ "validation_ids": [item.validation_id for item in result_list],
|
|
|
+ "errors": [item.error for item in result_list],
|
|
|
+ }
|
|
|
+ operation.completed_at = utc_now()
|
|
|
+ operation.updated_at = operation.completed_at
|
|
|
+ return {
|
|
|
+ "operation_id": operation_id,
|
|
|
+ "status": final_status.value,
|
|
|
+ "result_ref": operation.result_ref,
|
|
|
+ "error": error,
|
|
|
+ }
|
|
|
+
|
|
|
+ await self._mutate(
|
|
|
+ root_trace_id,
|
|
|
+ "operation_finished",
|
|
|
+ finish_state,
|
|
|
+ operation_id=operation_id,
|
|
|
+ )
|
|
|
+ return await self.get(root_trace_id, operation_id)
|
|
|
+
|
|
|
+
|
|
|
+def _recover_operation_state(
|
|
|
+ ledger: TaskLedger,
|
|
|
+ *,
|
|
|
+ operation_id: str,
|
|
|
+) -> Dict[str, Any]:
|
|
|
+ operation = ledger.operations[operation_id]
|
|
|
+ if operation.status in {
|
|
|
+ OperationStatus.COMPLETED,
|
|
|
+ OperationStatus.FAILED,
|
|
|
+ OperationStatus.STOPPED,
|
|
|
+ }:
|
|
|
+ return {"operation_id": operation_id, "status": operation.status.value}
|
|
|
+ expired = deadline_expired(operation.deadline_at)
|
|
|
+ if operation.status == OperationStatus.PENDING and not expired:
|
|
|
+ return {"operation_id": operation_id, "status": operation.status.value}
|
|
|
+
|
|
|
+ operation.status = OperationStatus.FAILED if expired else OperationStatus.STOPPED
|
|
|
+ operation.error = (
|
|
|
+ "Operation deadline exceeded during recovery"
|
|
|
+ if expired
|
|
|
+ else "Runtime restart requires explicit resume or Planner replan"
|
|
|
+ )
|
|
|
+ operation.completed_at = utc_now()
|
|
|
+ operation.updated_at = operation.completed_at
|
|
|
+ _recover_attempts(ledger, operation_id, operation.error, expired)
|
|
|
+ _recover_validations(ledger, operation_id, operation.error, expired)
|
|
|
+ if expired:
|
|
|
+ _mark_expired_tasks(ledger, operation.task_ids)
|
|
|
+ return {"operation_id": operation_id, "status": operation.status.value}
|
|
|
+
|
|
|
+
|
|
|
+def _recover_attempts(
|
|
|
+ ledger: TaskLedger,
|
|
|
+ operation_id: str,
|
|
|
+ error: str,
|
|
|
+ expired: bool,
|
|
|
+) -> None:
|
|
|
+ for attempt in ledger.attempts.values():
|
|
|
+ if not _attempt_needs_recovery(attempt, operation_id, expired):
|
|
|
+ continue
|
|
|
+ attempt.status = AttemptStatus.EXPIRED if expired else AttemptStatus.STOPPED
|
|
|
+ attempt.completed_at = utc_now()
|
|
|
+ attempt.error = attempt.error or error
|
|
|
+ attempt.execution_stats = _failure_stats(
|
|
|
+ attempt.execution_stats,
|
|
|
+ FailureCode.TIMEOUT if expired else FailureCode.INTERRUPTED,
|
|
|
+ )
|
|
|
+ task = ledger.tasks[attempt.task_id]
|
|
|
+ if task.status == TaskStatus.RUNNING:
|
|
|
+ task.status = TaskStatus.NEEDS_REPLAN
|
|
|
+
|
|
|
+
|
|
|
+def _attempt_needs_recovery(
|
|
|
+ attempt: TaskAttempt,
|
|
|
+ operation_id: str,
|
|
|
+ expired: bool,
|
|
|
+) -> bool:
|
|
|
+ return (
|
|
|
+ attempt.operation_id == operation_id
|
|
|
+ and attempt.status == AttemptStatus.RUNNING
|
|
|
+ and (expired or attempt.started_at is not None)
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _recover_validations(
|
|
|
+ ledger: TaskLedger,
|
|
|
+ operation_id: str,
|
|
|
+ error: str,
|
|
|
+ expired: bool,
|
|
|
+) -> None:
|
|
|
+ for validation in ledger.validations.values():
|
|
|
+ if not _validation_needs_recovery(validation, operation_id, expired):
|
|
|
+ continue
|
|
|
+ validation.status = (
|
|
|
+ ValidationRunStatus.EXPIRED if expired else ValidationRunStatus.STOPPED
|
|
|
+ )
|
|
|
+ validation.completed_at = utc_now()
|
|
|
+ validation.error = validation.error or error
|
|
|
+ validation.execution_stats = _failure_stats(
|
|
|
+ validation.execution_stats,
|
|
|
+ FailureCode.TIMEOUT if expired else FailureCode.INTERRUPTED,
|
|
|
+ )
|
|
|
+ task = ledger.tasks[validation.task_id]
|
|
|
+ if task.status == TaskStatus.VALIDATING:
|
|
|
+ task.status = TaskStatus.NEEDS_REPLAN
|
|
|
+
|
|
|
+
|
|
|
+def _validation_needs_recovery(
|
|
|
+ validation: ValidationReport,
|
|
|
+ operation_id: str,
|
|
|
+ expired: bool,
|
|
|
+) -> bool:
|
|
|
+ return (
|
|
|
+ validation.operation_id == operation_id
|
|
|
+ and validation.status == ValidationRunStatus.RUNNING
|
|
|
+ and (expired or validation.started_at is not None)
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _mark_expired_tasks(ledger: TaskLedger, task_ids: Sequence[str]) -> None:
|
|
|
+ active = {
|
|
|
+ TaskStatus.RUNNING,
|
|
|
+ TaskStatus.AWAITING_VALIDATION,
|
|
|
+ TaskStatus.VALIDATING,
|
|
|
+ }
|
|
|
+ for task_id in task_ids:
|
|
|
+ task = ledger.tasks[task_id]
|
|
|
+ if task.status in active:
|
|
|
+ task.status = TaskStatus.NEEDS_REPLAN
|
|
|
+ task.updated_at = utc_now()
|
|
|
+
|
|
|
+
|
|
|
+def _normalize_start_request(
|
|
|
+ kind: OperationKind | str,
|
|
|
+ task_ids: Optional[Sequence[str]],
|
|
|
+ task_id: Optional[str],
|
|
|
+ worker_presets: Optional[Sequence[str]],
|
|
|
+ attempt_id: Optional[str],
|
|
|
+ deadline_at: Optional[str],
|
|
|
+ default_worker_preset: str,
|
|
|
+) -> tuple[OperationKind, List[str], List[str], Dict[str, Any]]:
|
|
|
+ operation_kind = OperationKind(kind)
|
|
|
+ requested_task_ids = list(task_ids or ([] if task_id is None else [task_id]))
|
|
|
+ presets = list(worker_presets or [])
|
|
|
+ _validate_start_request(operation_kind, requested_task_ids, presets, attempt_id)
|
|
|
+ presets = _normalize_worker_presets(
|
|
|
+ operation_kind,
|
|
|
+ requested_task_ids,
|
|
|
+ presets,
|
|
|
+ default_worker_preset,
|
|
|
+ )
|
|
|
+ normalized_deadline = parse_deadline(deadline_at)
|
|
|
+ request = {
|
|
|
+ "kind": operation_kind.value,
|
|
|
+ "task_ids": requested_task_ids,
|
|
|
+ "worker_presets": presets,
|
|
|
+ "attempt_id": attempt_id,
|
|
|
+ "deadline_at": normalized_deadline,
|
|
|
+ }
|
|
|
+ return operation_kind, requested_task_ids, presets, request
|
|
|
+
|
|
|
+
|
|
|
+def _validate_start_request(
|
|
|
+ operation_kind: OperationKind,
|
|
|
+ task_ids: Sequence[str],
|
|
|
+ presets: Sequence[str],
|
|
|
+ attempt_id: Optional[str],
|
|
|
+) -> None:
|
|
|
+ if operation_kind == OperationKind.DISPATCH and not task_ids:
|
|
|
+ raise ValueError("Dispatch requires at least one task")
|
|
|
+ if len(set(task_ids)) != len(task_ids):
|
|
|
+ raise TaskConflict("The same task cannot be dispatched twice in one call")
|
|
|
+ if operation_kind == OperationKind.REVALIDATE:
|
|
|
+ if len(task_ids) != 1 or not attempt_id:
|
|
|
+ raise ValueError("Revalidate requires exactly one task_id and attempt_id")
|
|
|
+ if presets:
|
|
|
+ raise ValueError("worker_presets are not valid for revalidation")
|
|
|
+ if presets and len(presets) not in (1, len(task_ids)):
|
|
|
+ raise ValueError("worker_presets must have length 1 or match task_ids")
|
|
|
+
|
|
|
+
|
|
|
+def _normalize_worker_presets(
|
|
|
+ operation_kind: OperationKind,
|
|
|
+ task_ids: Sequence[str],
|
|
|
+ presets: List[str],
|
|
|
+ default_worker_preset: str,
|
|
|
+) -> List[str]:
|
|
|
+ if len(presets) == 1:
|
|
|
+ presets *= len(task_ids)
|
|
|
+ if operation_kind == OperationKind.DISPATCH and not presets:
|
|
|
+ return [default_worker_preset] * len(task_ids)
|
|
|
+ return presets
|
|
|
+
|
|
|
+
|
|
|
+def _create_operation(
|
|
|
+ ledger: TaskLedger,
|
|
|
+ *,
|
|
|
+ operation_id: str,
|
|
|
+ root_trace_id: str,
|
|
|
+ operation_kind: OperationKind,
|
|
|
+ request: Dict[str, Any],
|
|
|
+ fingerprint: str,
|
|
|
+ task_ids: Sequence[str],
|
|
|
+) -> Dict[str, Any]:
|
|
|
+ for task_id in task_ids:
|
|
|
+ if task_id not in ledger.tasks:
|
|
|
+ raise ValueError(f"Task not found: {task_id}")
|
|
|
+ ledger.operations[operation_id] = BackgroundOperation(
|
|
|
+ operation_id=operation_id,
|
|
|
+ root_trace_id=root_trace_id,
|
|
|
+ kind=operation_kind,
|
|
|
+ request=request,
|
|
|
+ request_fingerprint=fingerprint,
|
|
|
+ task_ids=list(task_ids),
|
|
|
+ deadline_at=request["deadline_at"],
|
|
|
+ )
|
|
|
+ return {"operation_id": operation_id}
|
|
|
+
|
|
|
+
|
|
|
+def parse_deadline(value: Optional[str]) -> Optional[str]:
|
|
|
+ if value is None:
|
|
|
+ return None
|
|
|
+ if not isinstance(value, str) or not value.strip():
|
|
|
+ raise ValueError("deadline_at must be an ISO-8601 timestamp")
|
|
|
+ try:
|
|
|
+ parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00"))
|
|
|
+ except ValueError as exc:
|
|
|
+ raise ValueError("deadline_at must be an ISO-8601 timestamp") from exc
|
|
|
+ if parsed.tzinfo is None or parsed.utcoffset() is None:
|
|
|
+ raise ValueError("deadline_at must include a timezone")
|
|
|
+ return parsed.astimezone(timezone.utc).isoformat()
|
|
|
+
|
|
|
+
|
|
|
+def deadline_expired(deadline: Optional[str]) -> bool:
|
|
|
+ return deadline is not None and deadline_seconds(deadline) <= 0
|
|
|
+
|
|
|
+
|
|
|
+def deadline_seconds(deadline: str) -> float:
|
|
|
+ parsed = datetime.fromisoformat(deadline.replace("Z", "+00:00"))
|
|
|
+ return (parsed.astimezone(timezone.utc) - datetime.now(timezone.utc)).total_seconds()
|
|
|
+
|
|
|
+
|
|
|
+def stage_timeout(
|
|
|
+ deadline: Optional[str],
|
|
|
+ configured_timeout: Optional[float],
|
|
|
+) -> Optional[float]:
|
|
|
+ if deadline is None:
|
|
|
+ return configured_timeout
|
|
|
+ remaining = max(0.0, deadline_seconds(deadline))
|
|
|
+ if configured_timeout is None:
|
|
|
+ return remaining
|
|
|
+ return min(configured_timeout, remaining)
|
|
|
+
|
|
|
+
|
|
|
+def _failure_stats(
|
|
|
+ stats: Optional[ExecutionStats],
|
|
|
+ fallback: FailureCode,
|
|
|
+) -> ExecutionStats:
|
|
|
+ if stats is None:
|
|
|
+ return ExecutionStats(failure_code=fallback)
|
|
|
+ if stats.failure_code is not None:
|
|
|
+ return stats
|
|
|
+ return replace(stats, failure_code=fallback)
|
|
|
+
|
|
|
+
|
|
|
+__all__ = ["OperationController", "stage_timeout"]
|
SamLee