77 lines
2.6 KiB
Python
77 lines
2.6 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
|