277 lines
10 KiB
Python
277 lines
10 KiB
Python
"""Task scheduler primitives."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from .errors import NotFoundError
|
|
from .capabilities import has_capabilities
|
|
from .models import WorkerClaim
|
|
from .store import TaskPlaneStore
|
|
|
|
|
|
class TaskScheduler:
|
|
def __init__(self, store: TaskPlaneStore):
|
|
self.store = store
|
|
|
|
def next_queued_task(self, worker_id: str):
|
|
worker = self.store.get_worker(worker_id)
|
|
tasks = sorted(
|
|
self.store.list_tasks(),
|
|
key=lambda task: (
|
|
-int(task.task_plan.metadata.get("priority", 0)),
|
|
task.created_at,
|
|
task.task_id,
|
|
),
|
|
)
|
|
for task in tasks:
|
|
if task.status not in {"CREATED", "QUEUED", "FAILED_RECOVERABLE"}:
|
|
continue
|
|
if task.task_plan.task_type == "worker.registry.audit":
|
|
continue
|
|
if task.status == "FAILED_RECOVERABLE" and not self.store.task_has_retryable_probe(
|
|
task.task_id
|
|
):
|
|
continue
|
|
if not has_capabilities(
|
|
worker.capabilities, task.task_plan.required_capabilities
|
|
):
|
|
continue
|
|
if any(
|
|
self.store.get_task(dependency).status != "COMPLETED"
|
|
for dependency in task.task_plan.dependencies
|
|
):
|
|
continue
|
|
if (
|
|
task.task_plan.data_classification == "credential-secret"
|
|
and worker.worker_type != "secure_file_transformer"
|
|
and "secure_file_transformer" not in worker.capabilities
|
|
):
|
|
continue
|
|
return task
|
|
return None
|
|
|
|
def schedule_for_worker(self, worker_id: str, lease_seconds: int = 15) -> WorkerClaim:
|
|
worker = self.store.get_worker(worker_id)
|
|
if worker.disabled:
|
|
from .errors import LockedError
|
|
|
|
raise LockedError("worker disabled")
|
|
if time.time() > worker.last_heartbeat + (
|
|
worker.heartbeat_interval_seconds * 3
|
|
):
|
|
from .errors import LockedError
|
|
|
|
raise LockedError("worker offline")
|
|
task = self.next_queued_task(worker_id)
|
|
if task is None:
|
|
raise NotFoundError("no queued tasks available")
|
|
claim = self.store.claim_task(
|
|
task.task_id,
|
|
worker_id,
|
|
worker.identity,
|
|
lease_seconds=lease_seconds,
|
|
expected_revision=task.revision,
|
|
operation_digest=task.task_plan.operation_digest(),
|
|
)
|
|
return claim
|
|
|
|
# R4.2 P2 durable multi-agent fleet scheduling public surface.
|
|
from .fleet import (
|
|
FleetConflict,
|
|
FleetDispatchDecision,
|
|
FleetError,
|
|
FleetScheduler,
|
|
FleetStateStore,
|
|
FleetWorkerPolicy,
|
|
)
|
|
# R4.2 P2-B fleet-aware pull scheduling.
|
|
#
|
|
# The original TaskScheduler remains the security oracle and claim-CAS implementation.
|
|
# Fleet ordering only controls the order presented to that oracle.
|
|
import inspect as _fleet_inspect
|
|
import threading as _fleet_threading
|
|
from collections.abc import Mapping as _FleetMapping
|
|
|
|
from .worker_registry import WorkerRegistry as _FleetWorkerRegistry
|
|
from .errors import NotFoundError as _FleetNotFoundError
|
|
|
|
_LegacyTaskScheduler = TaskScheduler
|
|
_LegacyScheduleParameters = _fleet_inspect.signature(
|
|
_LegacyTaskScheduler.schedule_for_worker
|
|
).parameters
|
|
_LegacyNextParameters = _fleet_inspect.signature(
|
|
_LegacyTaskScheduler.next_queued_task
|
|
).parameters
|
|
|
|
|
|
class _FleetOrderedStoreProxy:
|
|
def __init__(self, delegate, tasks):
|
|
self._delegate = delegate
|
|
self._tasks = tuple(tasks)
|
|
|
|
def list_tasks(self, *args, **kwargs):
|
|
return list(self._tasks)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._delegate, name)
|
|
|
|
|
|
def _fleet_claim_field(claim, name, default=None):
|
|
if isinstance(claim, _FleetMapping):
|
|
return claim.get(name, default)
|
|
return getattr(claim, name, default)
|
|
|
|
|
|
class TaskScheduler(_LegacyTaskScheduler):
|
|
"""Fleet-integrated pull scheduler preserving legacy security and claim semantics."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
store = getattr(self, "store", args[0] if args else None)
|
|
if store is None:
|
|
raise FleetError("FLEET_SCHEDULER_STORE_REQUIRED")
|
|
registry = getattr(self, "workers", None)
|
|
if registry is None:
|
|
registry = _FleetWorkerRegistry(store)
|
|
self.fleet_scheduler = FleetScheduler(store, registry)
|
|
self._fleet_local = _fleet_threading.local()
|
|
|
|
def _fleet_legacy_next(self, worker_id, tasks, *args, **kwargs):
|
|
proxy = _FleetOrderedStoreProxy(self.store, tasks)
|
|
legacy = _LegacyTaskScheduler(proxy)
|
|
legacy_kwargs = dict(kwargs)
|
|
if "actor_context" not in _LegacyNextParameters:
|
|
legacy_kwargs.pop("actor_context", None)
|
|
return legacy.next_queued_task(worker_id, *args, **legacy_kwargs)
|
|
|
|
def next_queued_task(self, worker_id, *args, **kwargs):
|
|
actor_context = kwargs.get(
|
|
"actor_context", getattr(self._fleet_local, "actor_context", None)
|
|
)
|
|
|
|
# Preserve legacy fail-closed worker checks before fleet ordering.
|
|
probe_kwargs = dict(kwargs)
|
|
if "actor_context" not in _LegacyNextParameters:
|
|
probe_kwargs.pop("actor_context", None)
|
|
legacy_probe = _LegacyTaskScheduler(self.store).next_queued_task(
|
|
worker_id, *args, **probe_kwargs
|
|
)
|
|
if legacy_probe is None:
|
|
return None
|
|
|
|
ranked = self.fleet_scheduler.rank_for_worker(
|
|
worker_id, actor_context=actor_context
|
|
)
|
|
# Preserve FleetScheduler fairness order. The legacy scheduler is used
|
|
# only as a fail-closed security oracle: present one fleet-ranked task
|
|
# at a time so its internal sort cannot reorder across scopes.
|
|
# Keep the legacy priority dimension ahead of Fleet fairness for
|
|
# compatibility with existing TaskPlane scheduling contracts; the
|
|
# Fleet order remains authoritative among equal-priority tasks.
|
|
def _legacy_priority(task):
|
|
plan = getattr(task, "task_plan", task)
|
|
metadata = getattr(plan, "metadata", {}) or {}
|
|
if not isinstance(metadata, _FleetMapping):
|
|
return 0
|
|
try:
|
|
return int(metadata.get("priority", 0))
|
|
except (TypeError, ValueError):
|
|
return 0
|
|
|
|
ranked = sorted(ranked, key=lambda task: -_legacy_priority(task))
|
|
for ranked_candidate in ranked:
|
|
try:
|
|
candidate = self._fleet_legacy_next(
|
|
worker_id, [ranked_candidate], *args, **kwargs
|
|
)
|
|
except _FleetNotFoundError:
|
|
continue
|
|
if candidate is None:
|
|
continue
|
|
|
|
ranked_id = getattr(
|
|
ranked_candidate,
|
|
"task_id",
|
|
getattr(getattr(ranked_candidate, "task_plan", None), "task_id", None),
|
|
)
|
|
candidate_id = getattr(
|
|
candidate,
|
|
"task_id",
|
|
getattr(getattr(candidate, "task_plan", None), "task_id", None),
|
|
)
|
|
if candidate_id != ranked_id:
|
|
raise FleetConflict("FLEET_LEGACY_ORACLE_ORDER_VIOLATION")
|
|
|
|
try:
|
|
decision = self.fleet_scheduler.reserve_for_worker(
|
|
worker_id, candidate
|
|
)
|
|
except FleetConflict:
|
|
continue
|
|
self._fleet_local.pending = decision
|
|
return candidate
|
|
return None
|
|
|
|
def schedule_for_worker(self, worker_id, *args, **kwargs):
|
|
actor_context = kwargs.get("actor_context")
|
|
legacy_kwargs = dict(kwargs)
|
|
if "actor_context" not in _LegacyScheduleParameters:
|
|
legacy_kwargs.pop("actor_context", None)
|
|
self._fleet_local.actor_context = actor_context
|
|
self._fleet_local.pending = None
|
|
try:
|
|
claim = _LegacyTaskScheduler.schedule_for_worker(
|
|
self, worker_id, *args, **legacy_kwargs
|
|
)
|
|
decision = getattr(self._fleet_local, "pending", None)
|
|
if claim is None:
|
|
if decision is not None:
|
|
self.fleet_scheduler.state.release(
|
|
decision.task_id, state="RELEASED"
|
|
)
|
|
return None
|
|
if decision is not None:
|
|
claim_id = _fleet_claim_field(claim, "claim_id")
|
|
claim_task_id = _fleet_claim_field(claim, "task_id", decision.task_id)
|
|
if not claim_id:
|
|
self.fleet_scheduler.state.release(
|
|
decision.task_id, state="RELEASED"
|
|
)
|
|
raise FleetError("FLEET_CLAIM_ID_REQUIRED")
|
|
if claim_task_id != decision.task_id:
|
|
self.fleet_scheduler.state.release(
|
|
decision.task_id, state="RELEASED"
|
|
)
|
|
raise FleetConflict("FLEET_CLAIM_TASK_MISMATCH")
|
|
self.fleet_scheduler.state.bind_claim(
|
|
decision.reservation_id, str(claim_id)
|
|
)
|
|
return claim
|
|
except Exception:
|
|
decision = getattr(self._fleet_local, "pending", None)
|
|
if decision is not None:
|
|
self.fleet_scheduler.state.release(
|
|
decision.task_id, state="RELEASED"
|
|
)
|
|
raise
|
|
finally:
|
|
self._fleet_local.pending = None
|
|
self._fleet_local.actor_context = None
|
|
|
|
def record_worker_failure(self, worker_id, failure_class, **kwargs):
|
|
return self.fleet_scheduler.state.record_worker_failure(
|
|
worker_id, failure_class, **kwargs
|
|
)
|
|
|
|
def record_worker_success(self, worker_id, **kwargs):
|
|
return self.fleet_scheduler.state.record_worker_success(
|
|
worker_id, **kwargs
|
|
)
|
|
|
|
def quarantine_task_worker(
|
|
self, task_id, worker_id, failure_class, **kwargs
|
|
):
|
|
return self.fleet_scheduler.state.record_task_worker_failure(
|
|
task_id, worker_id, failure_class, **kwargs
|
|
)
|