912 lines
34 KiB
Python
912 lines
34 KiB
Python
"""Durable multi-agent fleet scheduling for TaskPlane.
|
|
|
|
The fleet layer decides *which* eligible worker should receive *which* queued task.
|
|
It never replaces TaskPlane claim/lease ownership: callers must still perform the
|
|
authoritative TaskPlane claim CAS after a reservation is returned.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from typing import Any, Iterable, Mapping, Sequence
|
|
|
|
DEFAULT_RESERVATION_SECONDS = 30
|
|
DEFAULT_CIRCUIT_THRESHOLD = 3
|
|
DEFAULT_CIRCUIT_SECONDS = 30
|
|
MAX_CIRCUIT_SECONDS = 15 * 60
|
|
MAX_SLOTS = 128
|
|
|
|
|
|
class FleetError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class FleetConflict(FleetError):
|
|
pass
|
|
|
|
|
|
def _canonical(value: object) -> bytes:
|
|
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
|
|
|
|
|
def _sha(value: object) -> str:
|
|
return hashlib.sha256(_canonical(value)).hexdigest()
|
|
|
|
|
|
def _text(value: object, default: str = "") -> str:
|
|
if value is None:
|
|
return default
|
|
return str(value).strip() or default
|
|
|
|
|
|
def _int(value: object, default: int, minimum: int, maximum: int) -> int:
|
|
if value is None:
|
|
return default
|
|
if isinstance(value, bool):
|
|
raise FleetError("FLEET_INTEGER_INVALID")
|
|
try:
|
|
result = int(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise FleetError("FLEET_INTEGER_INVALID") from exc
|
|
if result < minimum or result > maximum:
|
|
raise FleetError("FLEET_INTEGER_OUT_OF_RANGE")
|
|
return result
|
|
|
|
|
|
def _float(value: object, default: float, minimum: float, maximum: float) -> float:
|
|
if value is None:
|
|
return default
|
|
try:
|
|
result = float(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise FleetError("FLEET_NUMBER_INVALID") from exc
|
|
if not math.isfinite(result) or result < minimum or result > maximum:
|
|
raise FleetError("FLEET_NUMBER_OUT_OF_RANGE")
|
|
return result
|
|
|
|
|
|
def _mapping(value: object) -> Mapping[str, Any]:
|
|
if isinstance(value, Mapping):
|
|
return value
|
|
if hasattr(value, "to_dict"):
|
|
data = value.to_dict()
|
|
if isinstance(data, Mapping):
|
|
return data
|
|
if hasattr(value, "__dict__"):
|
|
return vars(value)
|
|
return {}
|
|
|
|
|
|
def _attr(value: object, name: str, default: Any = None) -> Any:
|
|
if isinstance(value, Mapping):
|
|
return value.get(name, default)
|
|
return getattr(value, name, default)
|
|
|
|
|
|
def _task_plan(task: object) -> object:
|
|
return _attr(task, "task_plan", _attr(task, "plan", task))
|
|
|
|
|
|
def _task_id(task: object) -> str:
|
|
return _text(_attr(task, "task_id", _attr(_task_plan(task), "task_id", "")))
|
|
|
|
|
|
def _task_status(task: object) -> str:
|
|
return _text(_attr(task, "status", "")).upper()
|
|
|
|
|
|
def _task_revision(task: object) -> int:
|
|
return _int(_attr(task, "revision", 0), 0, 0, 2**31 - 1)
|
|
|
|
|
|
def _task_created_at(task: object) -> float:
|
|
value = _attr(task, "created_at", _attr(task, "updated_at", 0.0))
|
|
try:
|
|
return float(value or 0.0)
|
|
except (TypeError, ValueError):
|
|
return 0.0
|
|
|
|
|
|
def _task_required_capabilities(task: object) -> frozenset[str]:
|
|
plan = _task_plan(task)
|
|
value = _attr(plan, "required_capabilities", ())
|
|
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
|
|
return frozenset()
|
|
return frozenset(_text(item) for item in value if _text(item))
|
|
|
|
|
|
def _task_metadata(task: object) -> Mapping[str, Any]:
|
|
plan = _task_plan(task)
|
|
value = _attr(plan, "metadata", {})
|
|
return value if isinstance(value, Mapping) else {}
|
|
|
|
|
|
def _task_scope_key(task: object) -> str:
|
|
plan = _task_plan(task)
|
|
scope_type = _text(_attr(plan, "scope_type", "legacy"), "legacy")
|
|
tenant = _text(_attr(plan, "tenant_id", ""))
|
|
project = _text(_attr(plan, "project_id", ""))
|
|
user = _text(_attr(plan, "user_id", ""))
|
|
if scope_type == "tenant":
|
|
return "tenant:" + tenant + "|project:" + (project or "*")
|
|
if scope_type == "system":
|
|
return "system"
|
|
return "legacy:" + (user or "*")
|
|
|
|
|
|
def _worker_id(worker: object) -> str:
|
|
return _text(_attr(worker, "worker_id", ""))
|
|
|
|
|
|
def _worker_capabilities(worker: object) -> frozenset[str]:
|
|
value = _attr(worker, "capabilities", ())
|
|
if isinstance(value, Mapping):
|
|
value = value.keys()
|
|
if not isinstance(value, Iterable) or isinstance(value, (str, bytes)):
|
|
return frozenset()
|
|
return frozenset(_text(item) for item in value if _text(item))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FleetWorkerPolicy:
|
|
worker_id: str
|
|
max_slots: int = 1
|
|
failure_domain: str = "default"
|
|
weight: float = 1.0
|
|
labels: Mapping[str, str] | None = None
|
|
revision: int = 1
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FleetDispatchDecision:
|
|
reservation_id: str
|
|
task_id: str
|
|
worker_id: str
|
|
task_revision: int
|
|
scope_key: str
|
|
failure_domain: str
|
|
reserved_until: float
|
|
score: tuple[float, float, float, str]
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"reservation_id": self.reservation_id,
|
|
"task_id": self.task_id,
|
|
"worker_id": self.worker_id,
|
|
"task_revision": self.task_revision,
|
|
"scope_key": self.scope_key,
|
|
"failure_domain": self.failure_domain,
|
|
"reserved_until": self.reserved_until,
|
|
"score": list(self.score),
|
|
}
|
|
|
|
|
|
class FleetStateStore:
|
|
"""Small durable fleet journal colocated in the existing TaskPlane SQLite database."""
|
|
|
|
def __init__(self, connection: sqlite3.Connection, lock: threading.RLock | None = None) -> None:
|
|
self.conn = connection
|
|
self.lock = lock or threading.RLock()
|
|
self._savepoint_counter = 0
|
|
self._ensure_schema()
|
|
|
|
@contextmanager
|
|
def _transaction(self):
|
|
with self.lock:
|
|
self._savepoint_counter += 1
|
|
name = f"fleet_{threading.get_ident()}_{self._savepoint_counter}"
|
|
self.conn.execute(f"SAVEPOINT {name}")
|
|
try:
|
|
yield
|
|
except Exception:
|
|
self.conn.execute(f"ROLLBACK TO {name}")
|
|
self.conn.execute(f"RELEASE {name}")
|
|
raise
|
|
else:
|
|
self.conn.execute(f"RELEASE {name}")
|
|
|
|
def _ensure_schema(self) -> None:
|
|
with self.lock:
|
|
self.conn.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS fleet_worker_capacity (
|
|
worker_id TEXT PRIMARY KEY,
|
|
max_slots INTEGER NOT NULL,
|
|
failure_domain TEXT NOT NULL,
|
|
weight REAL NOT NULL,
|
|
labels_json TEXT NOT NULL,
|
|
revision INTEGER NOT NULL,
|
|
updated_at REAL NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS fleet_scope_state (
|
|
scope_key TEXT PRIMARY KEY,
|
|
virtual_finish REAL NOT NULL,
|
|
dispatch_count INTEGER NOT NULL,
|
|
last_dispatched_at REAL,
|
|
revision INTEGER NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS fleet_dispatches (
|
|
reservation_id TEXT PRIMARY KEY,
|
|
task_id TEXT NOT NULL UNIQUE,
|
|
worker_id TEXT NOT NULL,
|
|
task_revision INTEGER NOT NULL,
|
|
scope_key TEXT NOT NULL,
|
|
failure_domain TEXT NOT NULL,
|
|
state TEXT NOT NULL,
|
|
reserved_until REAL NOT NULL,
|
|
claim_id TEXT,
|
|
created_at REAL NOT NULL,
|
|
updated_at REAL NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS ix_fleet_dispatch_worker_state
|
|
ON fleet_dispatches(worker_id, state, reserved_until);
|
|
CREATE TABLE IF NOT EXISTS fleet_worker_failures (
|
|
worker_id TEXT PRIMARY KEY,
|
|
consecutive_failures INTEGER NOT NULL,
|
|
circuit_open_until REAL,
|
|
last_failure_class TEXT,
|
|
revision INTEGER NOT NULL,
|
|
updated_at REAL NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS fleet_task_worker_failures (
|
|
task_id TEXT NOT NULL,
|
|
worker_id TEXT NOT NULL,
|
|
failure_count INTEGER NOT NULL,
|
|
blocked_until REAL,
|
|
last_failure_class TEXT,
|
|
revision INTEGER NOT NULL,
|
|
updated_at REAL NOT NULL,
|
|
PRIMARY KEY(task_id, worker_id)
|
|
);
|
|
"""
|
|
)
|
|
|
|
def configure_worker(
|
|
self,
|
|
worker_id: str,
|
|
*,
|
|
max_slots: int = 1,
|
|
failure_domain: str = "default",
|
|
weight: float = 1.0,
|
|
labels: Mapping[str, str] | None = None,
|
|
expected_revision: int | None = None,
|
|
now: float | None = None,
|
|
) -> FleetWorkerPolicy:
|
|
worker_id = _text(worker_id)
|
|
if not worker_id:
|
|
raise FleetError("FLEET_WORKER_ID_REQUIRED")
|
|
max_slots = _int(max_slots, 1, 1, MAX_SLOTS)
|
|
weight = _float(weight, 1.0, 0.05, 100.0)
|
|
failure_domain = _text(failure_domain, "default")
|
|
normalized_labels = {
|
|
_text(k): _text(v)
|
|
for k, v in dict(labels or {}).items()
|
|
if _text(k) and _text(v)
|
|
}
|
|
now = float(time.time() if now is None else now)
|
|
with self._transaction():
|
|
row = self.conn.execute(
|
|
"SELECT revision FROM fleet_worker_capacity WHERE worker_id=?",
|
|
(worker_id,),
|
|
).fetchone()
|
|
current = int(row[0]) if row else 0
|
|
if expected_revision is not None and current != int(expected_revision):
|
|
raise FleetConflict("FLEET_WORKER_POLICY_CAS_MISMATCH")
|
|
revision = current + 1
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO fleet_worker_capacity(
|
|
worker_id,max_slots,failure_domain,weight,labels_json,revision,updated_at
|
|
) VALUES (?,?,?,?,?,?,?)
|
|
ON CONFLICT(worker_id) DO UPDATE SET
|
|
max_slots=excluded.max_slots,
|
|
failure_domain=excluded.failure_domain,
|
|
weight=excluded.weight,
|
|
labels_json=excluded.labels_json,
|
|
revision=excluded.revision,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(
|
|
worker_id,
|
|
max_slots,
|
|
failure_domain,
|
|
weight,
|
|
json.dumps(normalized_labels, sort_keys=True, separators=(",", ":")),
|
|
revision,
|
|
now,
|
|
),
|
|
)
|
|
return FleetWorkerPolicy(
|
|
worker_id, max_slots, failure_domain, weight, normalized_labels, revision
|
|
)
|
|
|
|
def worker_policy(self, worker_id: str) -> FleetWorkerPolicy:
|
|
worker_id = _text(worker_id)
|
|
with self.lock:
|
|
row = self.conn.execute(
|
|
"""
|
|
SELECT worker_id,max_slots,failure_domain,weight,labels_json,revision
|
|
FROM fleet_worker_capacity WHERE worker_id=?
|
|
""",
|
|
(worker_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
return FleetWorkerPolicy(worker_id=worker_id)
|
|
return FleetWorkerPolicy(
|
|
worker_id=str(row[0]),
|
|
max_slots=int(row[1]),
|
|
failure_domain=str(row[2]),
|
|
weight=float(row[3]),
|
|
labels=json.loads(row[4] or "{}"),
|
|
revision=int(row[5]),
|
|
)
|
|
|
|
def _cleanup_expired(self, now: float) -> None:
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE fleet_dispatches
|
|
SET state='EXPIRED', updated_at=?
|
|
WHERE state='RESERVED' AND reserved_until<=?
|
|
""",
|
|
(now, now),
|
|
)
|
|
|
|
def reserved_slots(self, worker_id: str, now: float) -> int:
|
|
with self.lock:
|
|
self._cleanup_expired(now)
|
|
row = self.conn.execute(
|
|
"""
|
|
SELECT COUNT(*) FROM fleet_dispatches
|
|
WHERE worker_id=? AND state='RESERVED' AND reserved_until>?
|
|
""",
|
|
(worker_id, now),
|
|
).fetchone()
|
|
return int(row[0] if row else 0)
|
|
|
|
def scope_virtual_finish(self, scope_key: str) -> float:
|
|
with self.lock:
|
|
row = self.conn.execute(
|
|
"SELECT virtual_finish FROM fleet_scope_state WHERE scope_key=?",
|
|
(scope_key,),
|
|
).fetchone()
|
|
return float(row[0]) if row else 0.0
|
|
|
|
def reserve(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
worker_id: str,
|
|
task_revision: int,
|
|
scope_key: str,
|
|
failure_domain: str,
|
|
score: tuple[float, float, float, str],
|
|
reservation_seconds: int,
|
|
now: float,
|
|
scope_weight: float,
|
|
) -> FleetDispatchDecision:
|
|
reservation_seconds = _int(
|
|
reservation_seconds, DEFAULT_RESERVATION_SECONDS, 1, 300
|
|
)
|
|
reservation_id = "fleet-" + _sha(
|
|
{
|
|
"task_id": task_id,
|
|
"worker_id": worker_id,
|
|
"task_revision": task_revision,
|
|
"scope_key": scope_key,
|
|
}
|
|
)[:32]
|
|
reserved_until = now + reservation_seconds
|
|
with self._transaction():
|
|
self._cleanup_expired(now)
|
|
existing = self.conn.execute(
|
|
"""
|
|
SELECT reservation_id,worker_id,task_revision,scope_key,failure_domain,
|
|
reserved_until,state
|
|
FROM fleet_dispatches WHERE task_id=?
|
|
""",
|
|
(task_id,),
|
|
).fetchone()
|
|
# Capacity is rechecked inside the same durable reservation transaction.
|
|
# This closes the cross-thread/process race between eligibility scoring and reserve.
|
|
policy_row = self.conn.execute(
|
|
"SELECT max_slots FROM fleet_worker_capacity WHERE worker_id=?",
|
|
(worker_id,),
|
|
).fetchone()
|
|
max_slots = int(policy_row[0]) if policy_row else 1
|
|
if existing and str(existing[6]) in {"RESERVED", "CLAIMED"} and float(existing[5]) > now:
|
|
if str(existing[1]) != worker_id or int(existing[2]) != task_revision:
|
|
raise FleetConflict("FLEET_TASK_ALREADY_RESERVED")
|
|
return FleetDispatchDecision(
|
|
str(existing[0]), task_id, worker_id, task_revision, str(existing[3]),
|
|
str(existing[4]), float(existing[5]), score
|
|
)
|
|
|
|
claim_columns = {
|
|
str(row[1])
|
|
for row in self.conn.execute("PRAGMA table_info(claims)").fetchall()
|
|
}
|
|
active_claims = 0
|
|
if {"worker_id", "status"}.issubset(claim_columns):
|
|
sql = "SELECT COUNT(*) FROM claims WHERE worker_id=? AND status='ACTIVE'"
|
|
params: list[Any] = [worker_id]
|
|
if "lease_expires_at" in claim_columns:
|
|
sql += " AND lease_expires_at>?"
|
|
params.append(now)
|
|
row = self.conn.execute(sql, tuple(params)).fetchone()
|
|
active_claims = int(row[0] if row else 0)
|
|
row = self.conn.execute(
|
|
"""
|
|
SELECT COUNT(*) FROM fleet_dispatches
|
|
WHERE worker_id=? AND task_id<>?
|
|
AND state='RESERVED' AND reserved_until>?
|
|
""",
|
|
(worker_id, task_id, now),
|
|
).fetchone()
|
|
active_reservations = int(row[0] if row else 0)
|
|
if active_claims + active_reservations >= max_slots:
|
|
raise FleetConflict("FLEET_WORKER_CAPACITY_EXHAUSTED")
|
|
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO fleet_dispatches(
|
|
reservation_id,task_id,worker_id,task_revision,scope_key,failure_domain,
|
|
state,reserved_until,claim_id,created_at,updated_at
|
|
) VALUES (?,?,?,?,?,?,'RESERVED',?,NULL,?,?)
|
|
ON CONFLICT(task_id) DO UPDATE SET
|
|
reservation_id=excluded.reservation_id,
|
|
worker_id=excluded.worker_id,
|
|
task_revision=excluded.task_revision,
|
|
scope_key=excluded.scope_key,
|
|
failure_domain=excluded.failure_domain,
|
|
state='RESERVED',
|
|
reserved_until=excluded.reserved_until,
|
|
claim_id=NULL,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(
|
|
reservation_id, task_id, worker_id, task_revision, scope_key,
|
|
failure_domain, reserved_until, now, now,
|
|
),
|
|
)
|
|
row = self.conn.execute(
|
|
"SELECT virtual_finish,dispatch_count,revision FROM fleet_scope_state WHERE scope_key=?",
|
|
(scope_key,),
|
|
).fetchone()
|
|
current_finish = float(row[0]) if row else 0.0
|
|
dispatch_count = int(row[1]) if row else 0
|
|
revision = int(row[2]) if row else 0
|
|
next_finish = max(current_finish, 0.0) + 1.0 / max(scope_weight, 0.05)
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO fleet_scope_state(
|
|
scope_key,virtual_finish,dispatch_count,last_dispatched_at,revision
|
|
) VALUES (?,?,?,?,?)
|
|
ON CONFLICT(scope_key) DO UPDATE SET
|
|
virtual_finish=excluded.virtual_finish,
|
|
dispatch_count=excluded.dispatch_count,
|
|
last_dispatched_at=excluded.last_dispatched_at,
|
|
revision=excluded.revision
|
|
""",
|
|
(scope_key, next_finish, dispatch_count + 1, now, revision + 1),
|
|
)
|
|
return FleetDispatchDecision(
|
|
reservation_id, task_id, worker_id, task_revision, scope_key,
|
|
failure_domain, reserved_until, score
|
|
)
|
|
|
|
def bind_claim(self, reservation_id: str, claim_id: str, *, now: float | None = None) -> None:
|
|
now = float(time.time() if now is None else now)
|
|
with self._transaction():
|
|
row = self.conn.execute(
|
|
"SELECT state FROM fleet_dispatches WHERE reservation_id=?",
|
|
(reservation_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise FleetError("FLEET_RESERVATION_NOT_FOUND")
|
|
if str(row[0]) not in {"RESERVED", "CLAIMED"}:
|
|
raise FleetConflict("FLEET_RESERVATION_NOT_ACTIVE")
|
|
self.conn.execute(
|
|
"""
|
|
UPDATE fleet_dispatches
|
|
SET state='CLAIMED', claim_id=?, updated_at=?
|
|
WHERE reservation_id=?
|
|
""",
|
|
(claim_id, now, reservation_id),
|
|
)
|
|
|
|
def release(
|
|
self,
|
|
task_id: str,
|
|
*,
|
|
state: str = "RELEASED",
|
|
now: float | None = None,
|
|
) -> None:
|
|
if state not in {"RELEASED", "COMPLETED", "FAILED", "EXPIRED"}:
|
|
raise FleetError("FLEET_RESERVATION_STATE_INVALID")
|
|
now = float(time.time() if now is None else now)
|
|
with self._transaction():
|
|
self.conn.execute(
|
|
"UPDATE fleet_dispatches SET state=?,updated_at=? WHERE task_id=?",
|
|
(state, now, task_id),
|
|
)
|
|
|
|
def worker_circuit_open(self, worker_id: str, now: float) -> bool:
|
|
with self.lock:
|
|
row = self.conn.execute(
|
|
"SELECT circuit_open_until FROM fleet_worker_failures WHERE worker_id=?",
|
|
(worker_id,),
|
|
).fetchone()
|
|
return bool(row and row[0] is not None and float(row[0]) > now)
|
|
|
|
def task_worker_blocked(self, task_id: str, worker_id: str, now: float) -> bool:
|
|
with self.lock:
|
|
row = self.conn.execute(
|
|
"""
|
|
SELECT blocked_until FROM fleet_task_worker_failures
|
|
WHERE task_id=? AND worker_id=?
|
|
""",
|
|
(task_id, worker_id),
|
|
).fetchone()
|
|
return bool(row and row[0] is not None and float(row[0]) > now)
|
|
|
|
def record_worker_failure(
|
|
self,
|
|
worker_id: str,
|
|
failure_class: str,
|
|
*,
|
|
threshold: int = DEFAULT_CIRCUIT_THRESHOLD,
|
|
base_circuit_seconds: int = DEFAULT_CIRCUIT_SECONDS,
|
|
now: float | None = None,
|
|
) -> dict[str, Any]:
|
|
now = float(time.time() if now is None else now)
|
|
threshold = _int(threshold, DEFAULT_CIRCUIT_THRESHOLD, 1, 20)
|
|
base_circuit_seconds = _int(base_circuit_seconds, DEFAULT_CIRCUIT_SECONDS, 1, 3600)
|
|
with self._transaction():
|
|
row = self.conn.execute(
|
|
"""
|
|
SELECT consecutive_failures,revision FROM fleet_worker_failures
|
|
WHERE worker_id=?
|
|
""",
|
|
(worker_id,),
|
|
).fetchone()
|
|
failures = (int(row[0]) if row else 0) + 1
|
|
revision = (int(row[1]) if row else 0) + 1
|
|
open_until = None
|
|
if failures >= threshold:
|
|
exponent = min(failures - threshold, 8)
|
|
open_until = now + min(
|
|
MAX_CIRCUIT_SECONDS, base_circuit_seconds * (2**exponent)
|
|
)
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO fleet_worker_failures(
|
|
worker_id,consecutive_failures,circuit_open_until,last_failure_class,
|
|
revision,updated_at
|
|
) VALUES (?,?,?,?,?,?)
|
|
ON CONFLICT(worker_id) DO UPDATE SET
|
|
consecutive_failures=excluded.consecutive_failures,
|
|
circuit_open_until=excluded.circuit_open_until,
|
|
last_failure_class=excluded.last_failure_class,
|
|
revision=excluded.revision,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(worker_id, failures, open_until, _text(failure_class, "unknown"), revision, now),
|
|
)
|
|
return {
|
|
"worker_id": worker_id,
|
|
"consecutive_failures": failures,
|
|
"circuit_open_until": open_until,
|
|
"revision": revision,
|
|
}
|
|
|
|
def record_worker_success(self, worker_id: str, *, now: float | None = None) -> None:
|
|
now = float(time.time() if now is None else now)
|
|
with self._transaction():
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO fleet_worker_failures(
|
|
worker_id,consecutive_failures,circuit_open_until,last_failure_class,
|
|
revision,updated_at
|
|
) VALUES (?,0,NULL,NULL,1,?)
|
|
ON CONFLICT(worker_id) DO UPDATE SET
|
|
consecutive_failures=0,
|
|
circuit_open_until=NULL,
|
|
last_failure_class=NULL,
|
|
revision=revision+1,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(worker_id, now),
|
|
)
|
|
|
|
def record_task_worker_failure(
|
|
self,
|
|
task_id: str,
|
|
worker_id: str,
|
|
failure_class: str,
|
|
*,
|
|
quarantine_seconds: int = 60,
|
|
now: float | None = None,
|
|
) -> dict[str, Any]:
|
|
now = float(time.time() if now is None else now)
|
|
quarantine_seconds = _int(quarantine_seconds, 60, 1, 3600)
|
|
with self._transaction():
|
|
row = self.conn.execute(
|
|
"""
|
|
SELECT failure_count,revision FROM fleet_task_worker_failures
|
|
WHERE task_id=? AND worker_id=?
|
|
""",
|
|
(task_id, worker_id),
|
|
).fetchone()
|
|
count = (int(row[0]) if row else 0) + 1
|
|
revision = (int(row[1]) if row else 0) + 1
|
|
blocked_until = now + quarantine_seconds
|
|
self.conn.execute(
|
|
"""
|
|
INSERT INTO fleet_task_worker_failures(
|
|
task_id,worker_id,failure_count,blocked_until,last_failure_class,
|
|
revision,updated_at
|
|
) VALUES (?,?,?,?,?,?,?)
|
|
ON CONFLICT(task_id,worker_id) DO UPDATE SET
|
|
failure_count=excluded.failure_count,
|
|
blocked_until=excluded.blocked_until,
|
|
last_failure_class=excluded.last_failure_class,
|
|
revision=excluded.revision,
|
|
updated_at=excluded.updated_at
|
|
""",
|
|
(
|
|
task_id, worker_id, count, blocked_until,
|
|
_text(failure_class, "unknown"), revision, now,
|
|
),
|
|
)
|
|
return {
|
|
"task_id": task_id,
|
|
"worker_id": worker_id,
|
|
"failure_count": count,
|
|
"blocked_until": blocked_until,
|
|
"revision": revision,
|
|
}
|
|
|
|
|
|
class FleetScheduler:
|
|
"""Deterministic capacity/fairness scheduler over existing TaskPlane records."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: Any,
|
|
workers: Any,
|
|
*,
|
|
now_fn: Any = time.time,
|
|
reservation_seconds: int = DEFAULT_RESERVATION_SECONDS,
|
|
) -> None:
|
|
self.store = store
|
|
self.workers = workers
|
|
self.now_fn = now_fn
|
|
self.reservation_seconds = _int(
|
|
reservation_seconds, DEFAULT_RESERVATION_SECONDS, 1, 300
|
|
)
|
|
connection = getattr(store, "conn", None)
|
|
if not isinstance(connection, sqlite3.Connection):
|
|
raise FleetError("FLEET_TASKPLANE_SQLITE_CONNECTION_REQUIRED")
|
|
lock = getattr(store, "_lock", None)
|
|
self.state = FleetStateStore(connection, lock)
|
|
|
|
def configure_worker(self, worker_id: str, **kwargs: Any) -> FleetWorkerPolicy:
|
|
return self.state.configure_worker(worker_id, **kwargs)
|
|
|
|
def _online(self, worker_id: str, now: float) -> bool:
|
|
method = getattr(self.workers, "is_online", None)
|
|
if not callable(method):
|
|
raise FleetError("FLEET_WORKER_REGISTRY_ONLINE_CHECK_REQUIRED")
|
|
try:
|
|
return bool(method(worker_id, now=now))
|
|
except TypeError:
|
|
return bool(method(worker_id))
|
|
|
|
def _worker_record(self, worker_id: str) -> object:
|
|
getter = getattr(self.store, "get_worker", None)
|
|
if not callable(getter):
|
|
raise FleetError("FLEET_WORKER_LOOKUP_REQUIRED")
|
|
return getter(worker_id)
|
|
|
|
def _active_claims(self, worker_id: str, now: float) -> int:
|
|
# The TaskPlane store owns one SQLite connection in the current runtime.
|
|
# Fleet reads must share its lock just like fleet writes do, otherwise two
|
|
# concurrent pullers can issue overlapping operations on the same handle.
|
|
with self.state.lock:
|
|
columns = {
|
|
str(row[1])
|
|
for row in self.store.conn.execute("PRAGMA table_info(claims)").fetchall()
|
|
}
|
|
required = {"worker_id", "status"}
|
|
if not required.issubset(columns):
|
|
raise FleetError("FLEET_CLAIMS_SCHEMA_UNSUPPORTED")
|
|
sql = "SELECT COUNT(*) FROM claims WHERE worker_id=? AND status='ACTIVE'"
|
|
params: list[Any] = [worker_id]
|
|
if "lease_expires_at" in columns:
|
|
sql += " AND lease_expires_at>?"
|
|
params.append(now)
|
|
row = self.store.conn.execute(sql, tuple(params)).fetchone()
|
|
return int(row[0] if row else 0)
|
|
|
|
@staticmethod
|
|
def _labels_match(policy: FleetWorkerPolicy, task: object) -> bool:
|
|
metadata = _task_metadata(task)
|
|
affinity = metadata.get("fleet_affinity_labels") or {}
|
|
if affinity and not isinstance(affinity, Mapping):
|
|
return False
|
|
labels = dict(policy.labels or {})
|
|
for key, value in dict(affinity).items():
|
|
if labels.get(_text(key)) != _text(value):
|
|
return False
|
|
avoid = metadata.get("fleet_avoid_failure_domains") or []
|
|
if isinstance(avoid, Sequence) and not isinstance(avoid, (str, bytes)):
|
|
if policy.failure_domain in {_text(item) for item in avoid}:
|
|
return False
|
|
return True
|
|
|
|
def _eligible_worker(self, worker_id: str, task: object, now: float) -> tuple[bool, int, FleetWorkerPolicy]:
|
|
if not self._online(worker_id, now):
|
|
return False, 0, self.state.worker_policy(worker_id)
|
|
worker = self._worker_record(worker_id)
|
|
policy = self.state.worker_policy(worker_id)
|
|
if self.state.worker_circuit_open(worker_id, now):
|
|
return False, 0, policy
|
|
if self.state.task_worker_blocked(_task_id(task), worker_id, now):
|
|
return False, 0, policy
|
|
if not self._labels_match(policy, task):
|
|
return False, 0, policy
|
|
required = _task_required_capabilities(task)
|
|
capabilities = _worker_capabilities(worker)
|
|
if required and not required.issubset(capabilities):
|
|
return False, 0, policy
|
|
active = self._active_claims(worker_id, now)
|
|
reserved = self.state.reserved_slots(worker_id, now)
|
|
load = active + reserved
|
|
if load >= policy.max_slots:
|
|
return False, load, policy
|
|
return True, load, policy
|
|
|
|
@staticmethod
|
|
def _task_priority(task: object, now: float) -> float:
|
|
metadata = _task_metadata(task)
|
|
base = _int(metadata.get("fleet_priority"), 0, -100, 100)
|
|
created = _task_created_at(task)
|
|
waited = max(0.0, now - created) if created > 0 else 0.0
|
|
aging = min(100.0, waited / 60.0)
|
|
return float(base) + aging
|
|
|
|
def _task_score(self, task: object, now: float) -> tuple[float, float, float, str]:
|
|
scope = _task_scope_key(task)
|
|
virtual_finish = self.state.scope_virtual_finish(scope)
|
|
priority = self._task_priority(task, now)
|
|
return (virtual_finish, -priority, _task_created_at(task), _task_id(task))
|
|
|
|
def _tasks(self, actor_context: Any = None) -> list[object]:
|
|
lister = getattr(self.store, "list_tasks", None)
|
|
if not callable(lister):
|
|
raise FleetError("FLEET_TASK_LIST_REQUIRED")
|
|
try:
|
|
tasks = lister(actor_context)
|
|
except TypeError:
|
|
tasks = lister()
|
|
return [
|
|
task for task in tasks
|
|
if _task_status(task) in {"CREATED", "QUEUED"}
|
|
and _task_id(task)
|
|
]
|
|
|
|
def rank_for_worker(
|
|
self,
|
|
worker_id: str,
|
|
*,
|
|
actor_context: Any = None,
|
|
now: float | None = None,
|
|
) -> list[object]:
|
|
"""Return fleet-ranked tasks without reserving or claiming anything."""
|
|
now = float(self.now_fn() if now is None else now)
|
|
worker_id = _text(worker_id)
|
|
if not worker_id:
|
|
return []
|
|
ranked: list[object] = []
|
|
for task in self._tasks(actor_context):
|
|
eligible, _, _ = self._eligible_worker(worker_id, task, now)
|
|
if eligible:
|
|
ranked.append(task)
|
|
return sorted(ranked, key=lambda task: self._task_score(task, now))
|
|
|
|
def reserve_for_worker(
|
|
self,
|
|
worker_id: str,
|
|
task: object,
|
|
*,
|
|
now: float | None = None,
|
|
) -> FleetDispatchDecision:
|
|
"""Reserve exactly one already-ranked task, rechecking eligibility and capacity."""
|
|
now = float(self.now_fn() if now is None else now)
|
|
worker_id = _text(worker_id)
|
|
eligible, _, policy = self._eligible_worker(worker_id, task, now)
|
|
if not eligible:
|
|
raise FleetConflict("FLEET_WORKER_NOT_ELIGIBLE")
|
|
metadata = _task_metadata(task)
|
|
scope_weight = _float(metadata.get("fleet_scope_weight"), 1.0, 0.05, 100.0)
|
|
return self.state.reserve(
|
|
task_id=_task_id(task),
|
|
worker_id=worker_id,
|
|
task_revision=_task_revision(task),
|
|
scope_key=_task_scope_key(task),
|
|
failure_domain=policy.failure_domain,
|
|
score=self._task_score(task, now),
|
|
reservation_seconds=self.reservation_seconds,
|
|
now=now,
|
|
scope_weight=scope_weight,
|
|
)
|
|
|
|
def select(
|
|
self,
|
|
worker_ids: Sequence[str],
|
|
*,
|
|
actor_context: Any = None,
|
|
limit: int = 1,
|
|
now: float | None = None,
|
|
) -> list[FleetDispatchDecision]:
|
|
now = float(self.now_fn() if now is None else now)
|
|
limit = _int(limit, 1, 1, 100)
|
|
worker_ids = tuple(dict.fromkeys(_text(item) for item in worker_ids if _text(item)))
|
|
tasks = sorted(self._tasks(actor_context), key=lambda task: self._task_score(task, now))
|
|
decisions: list[FleetDispatchDecision] = []
|
|
local_load: dict[str, int] = {}
|
|
for task in tasks:
|
|
if len(decisions) >= limit:
|
|
break
|
|
candidates: list[tuple[int, float, str, FleetWorkerPolicy]] = []
|
|
for worker_id in worker_ids:
|
|
eligible, load, policy = self._eligible_worker(worker_id, task, now)
|
|
if not eligible:
|
|
continue
|
|
effective_load = local_load.get(worker_id, load)
|
|
if effective_load >= policy.max_slots:
|
|
continue
|
|
candidates.append(
|
|
(
|
|
effective_load,
|
|
-policy.weight,
|
|
worker_id,
|
|
policy,
|
|
)
|
|
)
|
|
if not candidates:
|
|
continue
|
|
candidates.sort(key=lambda item: (item[0], item[1], item[2]))
|
|
_, _, worker_id, policy = candidates[0]
|
|
score = self._task_score(task, now)
|
|
metadata = _task_metadata(task)
|
|
scope_weight = _float(metadata.get("fleet_scope_weight"), 1.0, 0.05, 100.0)
|
|
try:
|
|
decision = self.state.reserve(
|
|
task_id=_task_id(task),
|
|
worker_id=worker_id,
|
|
task_revision=_task_revision(task),
|
|
scope_key=_task_scope_key(task),
|
|
failure_domain=policy.failure_domain,
|
|
score=score,
|
|
reservation_seconds=self.reservation_seconds,
|
|
now=now,
|
|
scope_weight=scope_weight,
|
|
)
|
|
except FleetConflict:
|
|
continue
|
|
local_load[worker_id] = local_load.get(
|
|
worker_id,
|
|
self._active_claims(worker_id, now)
|
|
+ self.state.reserved_slots(worker_id, now)
|
|
- 1,
|
|
) + 1
|
|
decisions.append(decision)
|
|
return decisions
|