The compute proof of concept had never been run with the agent as umbra-poc. Running it that way found three defects, all the same shape: a check that passed because its mechanism did nothing. The demonstration never reached the agent at all. It held the pid it spawned, which under `sudo -n -u` is the sudo wrapper - sudo forks the agent before exec'ing. Signalling that killed the wrapper and left the worker running, renewing its lease and holding the job indefinitely, under a printed note saying "exactly like a power cut". An unprivileged process cannot signal another account's process at all, and kill(2) on a process group succeeds as soon as one member is signalled, so the wrapper's death read as the agent's. Every run leaked an agent. Lease recovery was a diagnosis rather than a mechanism. sweep_expired_leases documented itself as "called by the controller's own timer" and there was no timer; a killed node's job came back only when somebody noticed a stalled job and posted to an administrative route by hand. The proof of concept swept by hand and recorded the result as automatic. app.py now runs the loop for the life of the application, building the compute repository at startup so a restart mid-job still recovers it. The power cut is now the agent's own: os._exit(137) the instant it claims a job, before running it, because a process that does not own the agent cannot signal it. A clean stop goes through a stop_file the agent polls, for the same reason. Liveness is measured by finding the worker's own pid in /proc matched on its config path, sampled repeatedly and accumulated - the wrapper and the worker are not simultaneous, so a single sample finds only the wrapper. The node key travels by `sudo --preserve-env` rather than as VAR=value in argv, because argv is world-readable through /proc/<pid>/cmdline and the environment is not. The demonstration checks whether the key will survive before starting and falls back to a pipe, which node_agent.py --node-key-stdin accepts. Neither a file nor argv ever holds it. Eight tests, each verified to fail against a mutation of the code it covers: three for the sweeper (it must return an expired lease, must survive a failing round, must not invent a repository it has no database for) and five driving a real agent process against a stub controller (stops on the stop file, exits 137 holding an uncompleted job, takes the key from both routes and neither, reports an empty pipe as an empty pipe). 448 pass with a database, 421 without. The deployment database is untouched: users=2 orgs=3 projects=0 audit=43, no compute tables.
1474 lines
60 KiB
Python
1474 lines
60 KiB
Python
"""Compute substrate: nodes, devices, credentials, the model registry, and the queue.
|
|
|
|
The interesting part of this module is not the CRUD. It is four decisions in
|
|
SPEC §6.16-6.18 and §7, each of which exists because the alternative is a failure
|
|
that is silent when it happens:
|
|
|
|
1. **A node is ``online`` because it heartbeated.** A node that was inserted and
|
|
never authenticated is ``provisioning`` and is offered no work. (I24)
|
|
2. **Claiming is ``FOR UPDATE SKIP LOCKED``.** Without it two workers select the
|
|
same row and one overwrites the other's claim, producing duplicated work that
|
|
is individually plausible and collectively wrong. (I25)
|
|
3. **A lease that expires returns the job to the queue.** A worker that is
|
|
``kill -9``'d releases nothing, so the lease is what releases it. (I26)
|
|
4. **Eligibility and preference are different questions.** A CPU device is slow,
|
|
not incapable, so a model that fits runs there rather than not running at all.
|
|
A job nothing can run is ``blocked`` with a reason, not left looking like it
|
|
is progressing. (§7.3, §7.3.1)
|
|
|
|
Nothing here knows about project content. A node credential reaches this module
|
|
and no further: it claims work and reports a result, and it cannot read a
|
|
manuscript. That boundary is the point of ``node_credentials`` existing as
|
|
something narrower than an API key.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any, Mapping, Sequence
|
|
from uuid import UUID, uuid4
|
|
|
|
import psycopg
|
|
|
|
from .errors import ConflictError, InputError, NotFoundError, RepositoryError
|
|
from .security import PasswordService, generate_token
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
#: A node that has not reported within this long is treated as absent. Three
|
|
#: missed heartbeats at a 10s interval, so a single slow or dropped report does
|
|
#: not take a working node out of service.
|
|
HEARTBEAT_STALE_SECONDS = 45
|
|
|
|
#: Default lease length for a claimed job. A node renews it while it works, so
|
|
#: this only has to exceed the time between renewals plus a little slack - it is
|
|
#: not a timeout for the work itself. See ``renew_lease`` and §6.17.
|
|
DEFAULT_LEASE_SECONDS = 60
|
|
|
|
#: How long a job is held back after a failed attempt. Grows with ``attempt`` so
|
|
#: a job that keeps failing stops hammering whichever node is unlucky enough to
|
|
#: claim it. The manifest's defect list names the absence of backoff by name.
|
|
BACKOFF_BASE_SECONDS = 5
|
|
BACKOFF_CEILING_SECONDS = 600
|
|
|
|
#: Statuses, spelled once. Mirrors the CHECK on ``jobs.status`` in migration 0003.
|
|
QUEUED = "queued"
|
|
RUNNING = "running"
|
|
BLOCKED = "blocked"
|
|
PARTIAL = "partial"
|
|
COMPLETE = "complete"
|
|
FAILED = "failed"
|
|
CANCELLED = "cancelled"
|
|
|
|
TERMINAL_STATUSES = frozenset({COMPLETE, FAILED, CANCELLED, PARTIAL})
|
|
|
|
#: The scopes a node credential may carry. A node credential is deliberately
|
|
#: unable to read project content; a scope that could would make it a session
|
|
#: token with a different name.
|
|
NODE_SCOPES = frozenset(
|
|
{
|
|
"compute:read", # read its own node, its jobs, the model registry
|
|
"compute:heartbeat", # report liveness and device measurements
|
|
"compute:claim", # claim a job
|
|
"compute:complete", # report a result, failure, or children
|
|
}
|
|
)
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Device:
|
|
"""One compute device on a node.
|
|
|
|
A node has a *set* of these rather than a capability flag, so "CPU and/or
|
|
GPU" is expressible and a CPU-only virtual machine is an ordinary node.
|
|
"""
|
|
|
|
id: UUID
|
|
compute_node_id: UUID
|
|
kind: str
|
|
label: str
|
|
vram_total_mb: int
|
|
vram_free_mb: int
|
|
ram_total_mb: int
|
|
ram_free_mb: int
|
|
compute_units: float | None = None
|
|
#: Orders placements. Never an eligibility test: a CPU device with a factor
|
|
#: of 1 against a GPU's 40 is slow, not incapable.
|
|
throughput_factor: float = 1.0
|
|
supports_accelerators: tuple[str, ...] = ()
|
|
|
|
@property
|
|
def uses_unified_memory(self) -> bool:
|
|
"""Whether this device's memory *is* system memory.
|
|
|
|
``cpu`` always, and ``gpu_integrated`` because integrated and Apple
|
|
silicon graphics memory is unified - the same bytes the CPU addresses. A
|
|
``gpu_discrete`` card does not: its VRAM is separate, and a model that
|
|
will not fit there runs on the *CPU*, which is a different device with its
|
|
own entry in ``compute_devices``.
|
|
|
|
Getting this wrong in the permissive direction lets a discrete card that is
|
|
out of VRAM silently take work it will then run on the processor under the
|
|
GPU's name, and the console will show a fast placement for a slow job.
|
|
"""
|
|
|
|
return self.kind in ("cpu", "gpu_integrated")
|
|
|
|
@property
|
|
def is_busy(self) -> bool:
|
|
"""Whether this device is out of memory.
|
|
|
|
Separate from ``throughput_factor`` on purpose. A node can be 0% utilised
|
|
and unable to take another model, because a resident model holds its
|
|
memory whether or not it is generating.
|
|
"""
|
|
|
|
return self.vram_free_mb <= 0 and self.kind != "cpu"
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ComputeNode:
|
|
"""A machine, and whether it has proved it is there."""
|
|
|
|
id: UUID
|
|
organization_id: UUID
|
|
name: str
|
|
is_controller: bool
|
|
trust_domain: str
|
|
status: str
|
|
locality: str | None = None
|
|
transport: str = "https"
|
|
agent_version: str | None = None
|
|
agent_started_at: datetime | None = None
|
|
last_heartbeat_at: datetime | None = None
|
|
created_at: datetime | None = None
|
|
devices: tuple[Device, ...] = ()
|
|
|
|
@property
|
|
def is_schedulable(self) -> bool:
|
|
"""Whether work may be dispatched here.
|
|
|
|
Requires both the status *and* a device that is not exhausted. A node
|
|
marked ``online`` whose only GPU is full is still not schedulable, and
|
|
the difference is exactly the case that makes a naive scheduler place a
|
|
job that then fails at inference time.
|
|
"""
|
|
|
|
if self.status != "online":
|
|
return False
|
|
return any(not device.is_busy for device in self.devices)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ModelRequirement:
|
|
"""What a job needs, and the minimum it can get away with.
|
|
|
|
``vram_min_mb`` is an eligibility floor; ``vram_preferred_mb`` is a ranking
|
|
target. Collapsing them into one number is the bug: it either refuses a job
|
|
that would have completed in an hour, or starts one that crawls.
|
|
"""
|
|
|
|
vram_min_mb: int = 0
|
|
ram_min_mb: int = 0
|
|
models: tuple[str, ...] = ()
|
|
trust_domains: tuple[str, ...] = ()
|
|
accelerators: tuple[str, ...] = ()
|
|
max_runtime_s: int | None = None
|
|
|
|
@classmethod
|
|
def from_json(cls, raw: Mapping[str, Any] | None) -> ModelRequirement:
|
|
data = raw or {}
|
|
return cls(
|
|
vram_min_mb=int(data.get("vram_min_mb", data.get("min_vram_mb", 0)) or 0),
|
|
ram_min_mb=int(data.get("ram_min_mb", data.get("min_ram_mb", 0)) or 0),
|
|
models=tuple(data.get("models", ()) or ()),
|
|
trust_domains=tuple(data.get("trust_domains", ()) or ()),
|
|
accelerators=tuple(data.get("accelerators", ()) or ()),
|
|
max_runtime_s=data.get("max_runtime_s"),
|
|
)
|
|
|
|
def to_json(self) -> dict[str, Any]:
|
|
return {
|
|
"vram_min_mb": self.vram_min_mb,
|
|
"ram_min_mb": self.ram_min_mb,
|
|
"models": list(self.models),
|
|
"trust_domains": list(self.trust_domains),
|
|
"accelerators": list(self.accelerators),
|
|
"max_runtime_s": self.max_runtime_s,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Placement:
|
|
"""A node chosen for a job, and the reason.
|
|
|
|
The reason is stored rather than implied because a scheduler that silently
|
|
prefers a fast node cannot be debugged when it starves a slow one. A job
|
|
running on a CPU device says so.
|
|
"""
|
|
|
|
node: ComputeNode
|
|
device: Device
|
|
reason: str
|
|
#: The model that is already resident, if any. A placement onto a device
|
|
#: where the model is loaded costs nothing; one that must load it, or evict
|
|
#: something to load it, costs minutes of held VRAM and somebody's work.
|
|
resident: bool = False
|
|
would_evict: tuple[str, ...] = ()
|
|
estimated_seconds: int | None = None
|
|
|
|
def to_json(self) -> dict[str, Any]:
|
|
return {
|
|
"node_id": str(self.node.id),
|
|
"node_name": self.node.name,
|
|
"device_id": str(self.device.id),
|
|
"device_kind": self.device.kind,
|
|
"reason": self.reason,
|
|
"resident": self.resident,
|
|
"would_evict": list(self.would_evict),
|
|
"estimated_seconds": self.estimated_seconds,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Verdict:
|
|
"""Whether a device can run a job, and which resource carries it.
|
|
|
|
``via`` is ``vram`` or ``ram``. It is the difference between a job that takes
|
|
four minutes and one that takes forty-seven, and it is why the reason string
|
|
is not a constant.
|
|
"""
|
|
|
|
via: str
|
|
would_evict: tuple[str, ...] = ()
|
|
resident: bool = False
|
|
|
|
@property
|
|
def is_slow_path(self) -> bool:
|
|
"""Running from system memory rather than dedicated VRAM.
|
|
|
|
Slow is not impossible. This is the flag that keeps the two apart.
|
|
"""
|
|
|
|
return self.via == "ram"
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Job:
|
|
"""One unit of queued work."""
|
|
|
|
id: UUID
|
|
organization_id: UUID
|
|
kind: str
|
|
status: str
|
|
requested_by: str
|
|
priority: int = 100
|
|
requires: ModelRequirement = field(default_factory=ModelRequirement)
|
|
payload: Mapping[str, Any] = field(default_factory=dict)
|
|
project_id: UUID | None = None
|
|
lang: str | None = None
|
|
lease_owner: UUID | None = None
|
|
lease_expires_at: datetime | None = None
|
|
attempt: int = 0
|
|
max_attempts: int = 3
|
|
not_before: datetime | None = None
|
|
parent_job_id: UUID | None = None
|
|
idempotency_key: str | None = None
|
|
result: Mapping[str, Any] | None = None
|
|
error: str | None = None
|
|
created_at: datetime | None = None
|
|
started_at: datetime | None = None
|
|
finished_at: datetime | None = None
|
|
|
|
@property
|
|
def is_terminal(self) -> bool:
|
|
return self.status in TERMINAL_STATUSES
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ClaimedJob:
|
|
"""A job, its placement, and the credential that claimed it.
|
|
|
|
The plan travels with the job so the node knows whether it is being asked to
|
|
load a model or evict one *before* it accepts. A node that discovers the
|
|
eviction after starting has already wasted the time.
|
|
"""
|
|
|
|
job: Job
|
|
placement: Placement
|
|
lease_expires_at: datetime
|
|
|
|
|
|
class ComputeRepository:
|
|
"""PostgreSQL implementation of the compute substrate."""
|
|
|
|
def __init__(
|
|
self,
|
|
dsn: str,
|
|
*,
|
|
min_size: int = 1,
|
|
max_size: int = 4,
|
|
passwords: PasswordService | None = None,
|
|
) -> None:
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
#: Node credentials are argon2id-hashed with the *same* profile as user
|
|
#: passwords, and through the same service, rather than a second hashing
|
|
#: path with its own defaults. Two profiles for two credential kinds is
|
|
#: two things to review, and one of them will be wrong.
|
|
self._passwords = passwords or PasswordService()
|
|
self._pool = ConnectionPool(
|
|
dsn,
|
|
min_size=min_size,
|
|
max_size=max_size,
|
|
open=True,
|
|
kwargs={"row_factory": psycopg.rows.dict_row, "autocommit": False},
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self._pool.close()
|
|
|
|
# -- connections -------------------------------------------------------
|
|
|
|
def _connection(self):
|
|
return self._pool.connection()
|
|
|
|
@staticmethod
|
|
def _tx(connection):
|
|
return connection.transaction()
|
|
|
|
# -- nodes -------------------------------------------------------------
|
|
|
|
def create_node(
|
|
self,
|
|
organization_id: UUID,
|
|
*,
|
|
name: str,
|
|
is_controller: bool = False,
|
|
trust_domain: str = "local",
|
|
locality: str | None = None,
|
|
transport: str = "https",
|
|
) -> ComputeNode:
|
|
"""Insert a node in ``provisioning``.
|
|
|
|
Not ``online``. ``online`` is earned by a heartbeat, and a node that
|
|
looks available before its agent has ever spoken is a node that silently
|
|
swallows dispatched work. (I24)
|
|
"""
|
|
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.compute_nodes
|
|
(organization_id, name, is_controller, trust_domain, locality, transport, status)
|
|
VALUES (%s, %s, %s, %s, %s, %s, 'provisioning')
|
|
RETURNING *
|
|
""",
|
|
(organization_id, name, is_controller, trust_domain, locality, transport),
|
|
).fetchone()
|
|
return _node_from_row(row, devices=())
|
|
|
|
def get_node(self, node_id: UUID) -> ComputeNode | None:
|
|
with self._connection() as connection:
|
|
return self._load_node(connection, node_id)
|
|
|
|
def _load_node(self, connection, node_id: UUID) -> ComputeNode | None:
|
|
row = connection.execute(
|
|
"SELECT * FROM public.compute_nodes WHERE id = %s", (node_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
devices = connection.execute(
|
|
"SELECT * FROM public.compute_devices WHERE compute_node_id = %s ORDER BY label",
|
|
(node_id,),
|
|
).fetchall()
|
|
return _node_from_row(row, devices=tuple(_device_from_row(d) for d in devices))
|
|
|
|
def list_nodes(self, organization_id: UUID) -> list[ComputeNode]:
|
|
with self._connection() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM public.compute_nodes WHERE organization_id = %s ORDER BY name",
|
|
(organization_id,),
|
|
).fetchall()
|
|
return [
|
|
_node_from_row(
|
|
row,
|
|
devices=tuple(
|
|
_device_from_row(d)
|
|
for d in connection.execute(
|
|
"SELECT * FROM public.compute_devices WHERE compute_node_id = %s ORDER BY label",
|
|
(row["id"],),
|
|
).fetchall()
|
|
),
|
|
)
|
|
for row in rows
|
|
]
|
|
|
|
def set_node_status(self, node_id: UUID, status: str) -> ComputeNode:
|
|
"""An administrator's control: drain, retire, return to service.
|
|
|
|
Not the path to ``online``. Only :meth:`heartbeat` sets that, which is
|
|
what makes ``online`` mean "I am here" rather than "someone typed it".
|
|
"""
|
|
|
|
allowed = {"provisioning", "draining", "offline", "retired"}
|
|
if status not in allowed:
|
|
raise RepositoryError(
|
|
"invalid_node_status",
|
|
"Node status may only be set by an administrator to provisioning, "
|
|
"draining, offline, or retired. 'online' is earned by a heartbeat.",
|
|
)
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"UPDATE public.compute_nodes SET status = %s, updated_at = now() "
|
|
"WHERE id = %s RETURNING id",
|
|
(status, node_id),
|
|
).fetchone()
|
|
if row is None:
|
|
raise NotFoundError("compute_node_not_found")
|
|
node = self.get_node(node_id)
|
|
assert node is not None
|
|
return node
|
|
|
|
# -- devices -----------------------------------------------------------
|
|
|
|
def report_devices(self, node_id: UUID, devices: Sequence[Mapping[str, Any]]) -> list[Device]:
|
|
"""Replace a node's device measurements with what it just measured.
|
|
|
|
Upsert by ``(node, label)`` rather than by id, because the node's
|
|
device ids are its own and a restart with a different enumeration should
|
|
not collide with the previous run's rows.
|
|
"""
|
|
|
|
with self._connection() as connection, self._tx(connection):
|
|
for spec in devices:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO public.compute_devices
|
|
(compute_node_id, kind, label, vram_total_mb, vram_free_mb,
|
|
ram_total_mb, ram_free_mb, compute_units, throughput_factor,
|
|
supports_accelerators)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (compute_node_id, label) DO UPDATE SET
|
|
kind = EXCLUDED.kind,
|
|
vram_total_mb = EXCLUDED.vram_total_mb,
|
|
vram_free_mb = EXCLUDED.vram_free_mb,
|
|
ram_total_mb = EXCLUDED.ram_total_mb,
|
|
ram_free_mb = EXCLUDED.ram_free_mb,
|
|
compute_units = EXCLUDED.compute_units,
|
|
throughput_factor = EXCLUDED.throughput_factor,
|
|
supports_accelerators = EXCLUDED.supports_accelerators,
|
|
updated_at = now()
|
|
""",
|
|
(
|
|
node_id,
|
|
spec["kind"],
|
|
spec["label"],
|
|
int(spec.get("vram_total_mb", 0)),
|
|
int(spec.get("vram_free_mb", 0)),
|
|
int(spec.get("ram_total_mb", 0)),
|
|
int(spec.get("ram_free_mb", 0)),
|
|
spec.get("compute_units"),
|
|
float(spec.get("throughput_factor", 1.0)),
|
|
list(spec.get("supports_accelerators", ())),
|
|
),
|
|
)
|
|
node = self.get_node(node_id)
|
|
assert node is not None
|
|
return list(node.devices)
|
|
|
|
# -- credentials -------------------------------------------------------
|
|
|
|
def issue_node_credential(
|
|
self, node_id: UUID, *, agent_name: str, scopes: Sequence[str]
|
|
) -> tuple[str, str]:
|
|
"""Return ``(prefix, key)``. The key is shown once and never stored.
|
|
|
|
Scopes are validated against :data:`NODE_SCOPES` here rather than at
|
|
the call site, because a scope outside that set would be a node
|
|
credential that can read project content - which is the boundary this
|
|
table exists to hold.
|
|
"""
|
|
|
|
wanted = sorted(set(scopes))
|
|
if not wanted:
|
|
raise InputError("node_scope_required", "A node credential needs at least one scope.")
|
|
unknown = [scope for scope in wanted if scope not in NODE_SCOPES]
|
|
if unknown:
|
|
raise InputError(
|
|
"node_scope_unknown",
|
|
f"These are not node scopes: {', '.join(unknown)}. A node credential "
|
|
f"may only carry: {', '.join(sorted(NODE_SCOPES))}.",
|
|
)
|
|
key = generate_token(32)
|
|
prefix = key[:8]
|
|
with self._connection() as connection, self._tx(connection):
|
|
connection.execute(
|
|
"INSERT INTO public.node_credentials "
|
|
"(compute_node_id, key_prefix, key_hash, agent_name, scopes) "
|
|
"VALUES (%s, %s, %s, %s, %s)",
|
|
(node_id, prefix, self._passwords.hash_password(key), agent_name, wanted),
|
|
)
|
|
return prefix, key
|
|
|
|
def resolve_node_credential(self, key: str) -> tuple[UUID, UUID, list[str]] | None:
|
|
"""Turn a presented key into ``(node_id, organization_id, scopes)``.
|
|
|
|
The organization comes from the node, never from the request. That is
|
|
what makes a forged tenant identifier unnecessary rather than merely
|
|
rejected, and it is why a node needs no database credential at all.
|
|
"""
|
|
|
|
with self._connection() as connection:
|
|
row = connection.execute(
|
|
"""
|
|
SELECT c.compute_node_id, n.organization_id, c.scopes, c.key_hash, c.revoked_at
|
|
FROM public.node_credentials c
|
|
JOIN public.compute_nodes n ON n.id = c.compute_node_id
|
|
WHERE c.key_prefix = %s
|
|
""",
|
|
(key[:8],),
|
|
).fetchone()
|
|
if row is None or row["revoked_at"] is not None:
|
|
return None
|
|
if not self._passwords.verify_password(row["key_hash"], key):
|
|
return None
|
|
return row["compute_node_id"], row["organization_id"], list(row["scopes"])
|
|
|
|
def mark_credential_used(self, node_id: UUID) -> None:
|
|
with self._connection() as connection, self._tx(connection):
|
|
connection.execute(
|
|
"UPDATE public.node_credentials SET last_used_at = now() WHERE compute_node_id = %s",
|
|
(node_id,),
|
|
)
|
|
|
|
# -- heartbeat ---------------------------------------------------------
|
|
|
|
def heartbeat(
|
|
self,
|
|
node_id: UUID,
|
|
*,
|
|
agent_version: str | None = None,
|
|
measurements: Mapping[str, Any] | None = None,
|
|
resident_models: Sequence[str] = (),
|
|
active_jobs: int = 0,
|
|
devices: Sequence[Mapping[str, Any]] = (),
|
|
) -> ComputeNode:
|
|
"""Record liveness and **promote the node to online**.
|
|
|
|
The promotion happens here and nowhere else, which is what makes
|
|
``online`` mean "this node has spoken" rather than "this row exists".
|
|
(I24)
|
|
|
|
A node an administrator has set to ``draining`` or ``retired`` is not
|
|
promoted back by heartbeating. Otherwise a draining node that is not yet
|
|
rebooted would re-advertise itself for work the administrator is trying
|
|
to drain away from it.
|
|
"""
|
|
|
|
data = measurements or {}
|
|
with self._connection() as connection, self._tx(connection):
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO public.node_heartbeats
|
|
(compute_node_id, reported_at, cpu_busy_pct, ram_free_mb,
|
|
vram_free_mb, vram_total_mb, disk_free_mb, load_avg_1m,
|
|
resident_models, active_jobs, agent_version, raw)
|
|
VALUES (%s, now(), %s, %s, %s, %s, %s, %s, %s::jsonb, %s, %s, %s::jsonb)
|
|
ON CONFLICT (compute_node_id) DO UPDATE SET
|
|
reported_at = now(),
|
|
cpu_busy_pct = EXCLUDED.cpu_busy_pct,
|
|
ram_free_mb = EXCLUDED.ram_free_mb,
|
|
vram_free_mb = EXCLUDED.vram_free_mb,
|
|
vram_total_mb = EXCLUDED.vram_total_mb,
|
|
disk_free_mb = EXCLUDED.disk_free_mb,
|
|
load_avg_1m = EXCLUDED.load_avg_1m,
|
|
resident_models = EXCLUDED.resident_models,
|
|
active_jobs = EXCLUDED.active_jobs,
|
|
agent_version = EXCLUDED.agent_version,
|
|
raw = EXCLUDED.raw
|
|
""",
|
|
(
|
|
node_id,
|
|
data.get("cpu_busy_pct"),
|
|
data.get("ram_free_mb"),
|
|
data.get("vram_free_mb"),
|
|
data.get("vram_total_mb"),
|
|
data.get("disk_free_mb"),
|
|
data.get("load_avg_1m"),
|
|
json.dumps(list(resident_models)),
|
|
active_jobs,
|
|
agent_version,
|
|
json.dumps(dict(data)),
|
|
),
|
|
)
|
|
connection.execute(
|
|
"""
|
|
UPDATE public.compute_nodes
|
|
SET last_heartbeat_at = now(),
|
|
agent_version = %s,
|
|
agent_started_at = now(),
|
|
status = CASE WHEN status IN ('provisioning', 'offline') THEN 'online'
|
|
ELSE status END,
|
|
updated_at = now()
|
|
WHERE id = %s
|
|
""",
|
|
(agent_version, node_id),
|
|
)
|
|
if devices:
|
|
self.report_devices(node_id, devices)
|
|
node = self.get_node(node_id)
|
|
if node is None:
|
|
raise NotFoundError("compute_node_not_found")
|
|
return node
|
|
|
|
def get_heartbeat(self, node_id: UUID) -> dict[str, Any] | None:
|
|
with self._connection() as connection:
|
|
row = connection.execute(
|
|
"SELECT * FROM public.node_heartbeats WHERE compute_node_id = %s", (node_id,)
|
|
).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
# -- model registry ----------------------------------------------------
|
|
|
|
def register_model(
|
|
self,
|
|
slug: str,
|
|
*,
|
|
kind: str = "llm",
|
|
vram_min_mb: int = 0,
|
|
vram_preferred_mb: int = 0,
|
|
ram_min_mb: int = 0,
|
|
requires_accelerators: Sequence[str] = (),
|
|
reference_throughput: float = 1.0,
|
|
load_seconds: int = 0,
|
|
) -> dict[str, Any]:
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.model_registry
|
|
(slug, kind, vram_min_mb, vram_preferred_mb, ram_min_mb,
|
|
requires_accelerators, reference_throughput, load_seconds)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (slug) DO UPDATE SET
|
|
kind = EXCLUDED.kind,
|
|
vram_min_mb = EXCLUDED.vram_min_mb,
|
|
vram_preferred_mb = EXCLUDED.vram_preferred_mb,
|
|
ram_min_mb = EXCLUDED.ram_min_mb,
|
|
requires_accelerators = EXCLUDED.requires_accelerators,
|
|
reference_throughput = EXCLUDED.reference_throughput,
|
|
load_seconds = EXCLUDED.load_seconds
|
|
RETURNING *
|
|
""",
|
|
(
|
|
slug,
|
|
kind,
|
|
vram_min_mb,
|
|
vram_preferred_mb,
|
|
ram_min_mb,
|
|
list(requires_accelerators),
|
|
reference_throughput,
|
|
load_seconds,
|
|
),
|
|
).fetchone()
|
|
return dict(row)
|
|
|
|
def list_models(self) -> list[dict[str, Any]]:
|
|
with self._connection() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM public.model_registry ORDER BY slug"
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def set_node_model_state(
|
|
self, node_id: UUID, slug: str, state: str, *, device_id: UUID | None = None, vram_held_mb: int = 0
|
|
) -> None:
|
|
"""Record what is resident on a node.
|
|
|
|
This is the input to the residency half of placement scoring, and the
|
|
reason a node at 0% utilisation with 87% of its VRAM held is visible as
|
|
constrained rather than as idle.
|
|
"""
|
|
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"SELECT id FROM public.model_registry WHERE slug = %s", (slug,)
|
|
).fetchone()
|
|
if row is None:
|
|
raise NotFoundError("model_not_registered")
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO public.node_models
|
|
(compute_node_id, model_id, device_id, state, loaded_at, vram_held_mb)
|
|
VALUES (%s, %s, %s, %s, now(), %s)
|
|
ON CONFLICT (compute_node_id, model_id) DO UPDATE SET
|
|
device_id = EXCLUDED.device_id,
|
|
state = EXCLUDED.state,
|
|
loaded_at = CASE WHEN EXCLUDED.state = 'loaded' THEN now()
|
|
ELSE node_models.loaded_at END,
|
|
vram_held_mb = EXCLUDED.vram_held_mb,
|
|
updated_at = now()
|
|
""",
|
|
(node_id, row["id"], device_id, state, vram_held_mb),
|
|
)
|
|
|
|
def _resident_models(self, connection, node_id: UUID) -> dict[str, dict[str, Any]]:
|
|
rows = connection.execute(
|
|
"""
|
|
SELECT m.slug, nm.state, nm.device_id, nm.vram_held_mb
|
|
FROM public.node_models nm
|
|
JOIN public.model_registry m ON m.id = nm.model_id
|
|
WHERE nm.compute_node_id = %s AND nm.state = 'loaded'
|
|
""",
|
|
(node_id,),
|
|
).fetchall()
|
|
return {row["slug"]: dict(row) for row in rows}
|
|
|
|
# -- placement ---------------------------------------------------------
|
|
|
|
def place(
|
|
self, requirement: ModelRequirement, *, organization_id: UUID, prefer_node: UUID | None = None
|
|
) -> Placement | None:
|
|
"""Choose a node and a device, or return ``None``.
|
|
|
|
Two questions, answered separately.
|
|
|
|
**Eligibility** - can it run here at all? A device qualifies on free
|
|
memory and accelerator support. Nothing about speed.
|
|
|
|
**Preference** - how fast, and what does arriving cost? Scored by
|
|
throughput, discounted for having to load the model, and discounted
|
|
further for having to evict somebody to load it.
|
|
|
|
That split is the whole of §7.3. A device that is slow is eligible; a
|
|
device that is too small is not. Returning ``None`` for the former is the
|
|
bug, and it is the expensive one, because it turns a machine that would
|
|
have finished the work in an hour into a machine that never touches it.
|
|
"""
|
|
|
|
with self._connection() as connection:
|
|
node_rows = connection.execute(
|
|
"""
|
|
SELECT * FROM public.compute_nodes
|
|
WHERE organization_id = %s AND status = 'online'
|
|
ORDER BY name
|
|
""",
|
|
(organization_id,),
|
|
).fetchall()
|
|
models = {
|
|
row["slug"]: dict(row)
|
|
for row in connection.execute("SELECT * FROM public.model_registry").fetchall()
|
|
}
|
|
|
|
candidates: list[Placement] = []
|
|
for row in node_rows:
|
|
last_seen = row["last_heartbeat_at"]
|
|
if last_seen is None:
|
|
# Online but never heard from is a contradiction. Only the
|
|
# heartbeat path sets 'online', so this row is stale or was
|
|
# written by something else. Either way: absent.
|
|
continue
|
|
age = (_now() - now_aware(last_seen)).total_seconds()
|
|
if age > HEARTBEAT_STALE_SECONDS:
|
|
# Silent for too long. It is not a slow node, it is a node
|
|
# that has gone, and offering it work means the job is
|
|
# dispatched into the void and recovered by the lease sweep
|
|
# several minutes later instead of being offered to a machine
|
|
# that would have taken it immediately.
|
|
continue
|
|
node = _node_from_row(row, devices=())
|
|
devices = connection.execute(
|
|
"SELECT * FROM public.compute_devices WHERE compute_node_id = %s ORDER BY label",
|
|
(row["id"],),
|
|
).fetchall()
|
|
node = _node_from_row(row, devices=tuple(_device_from_row(d) for d in devices))
|
|
if requirement.trust_domains and node.trust_domain not in requirement.trust_domains:
|
|
continue
|
|
resident = self._resident_models(connection, node.id)
|
|
for device_row in devices:
|
|
device = _device_from_row(device_row)
|
|
verdict = self._eligible(device, requirement, models, resident)
|
|
if verdict is None:
|
|
continue
|
|
candidates.append(
|
|
Placement(
|
|
node=node,
|
|
device=device,
|
|
reason=_placement_reason(
|
|
device, prefer_node, node.id == prefer_node, verdict.is_slow_path
|
|
),
|
|
resident=verdict.resident,
|
|
would_evict=verdict.would_evict,
|
|
)
|
|
)
|
|
if not candidates:
|
|
return None
|
|
return max(candidates, key=_placement_score)
|
|
|
|
@staticmethod
|
|
def _eligible(
|
|
device: Device,
|
|
requirement: ModelRequirement,
|
|
models: Mapping[str, Mapping[str, Any]],
|
|
resident: Mapping[str, Mapping[str, Any]],
|
|
) -> "Verdict | None":
|
|
"""Decide whether this device can run the job at all, and how.
|
|
|
|
**A model that does not fit in VRAM can still run on system RAM.** That is
|
|
the rule this function got wrong first, and getting it wrong is expensive
|
|
in the direction that is hard to notice: a CPU device, or a machine with
|
|
unified graphics memory, has *zero* dedicated VRAM, so a VRAM-only test
|
|
disqualifies every one of them. The effect is that a machine which would
|
|
have finished the work in an hour is never offered it at all.
|
|
|
|
So the test is "VRAM **or** RAM", and the verdict records which one
|
|
qualified, because the difference between the two is the difference
|
|
between minutes and hours and belongs in the placement reason.
|
|
"""
|
|
|
|
if device.is_busy:
|
|
return None
|
|
|
|
ram_floor = requirement.ram_min_mb
|
|
for slug in requirement.models:
|
|
model = models.get(slug)
|
|
if model is None:
|
|
return None
|
|
if model["requires_accelerators"] and not set(model["requires_accelerators"]) <= set(
|
|
device.supports_accelerators
|
|
):
|
|
return None
|
|
ram_floor = max(ram_floor, int(model["ram_min_mb"] or 0))
|
|
if not model["ram_min_mb"] and model["vram_min_mb"]:
|
|
# A stated VRAM floor with no stated RAM floor means the model has
|
|
# not been measured for RAM yet. Its VRAM floor is a proxy for how
|
|
# big it is, and that size does not go away when it runs on the
|
|
# processor. Without this, a model needing 12 GB of VRAM would be
|
|
# offered to a machine with 2 GB of RAM, because "no RAM floor
|
|
# stated" would read as "any RAM will do".
|
|
ram_floor = max(ram_floor, int(model["vram_min_mb"]))
|
|
|
|
vram_floor = max(
|
|
[requirement.vram_min_mb, *[int(models[s]["vram_min_mb"] or 0) for s in requirement.models if s in models]]
|
|
or [0]
|
|
)
|
|
|
|
vram_ok = vram_floor == 0 or device.vram_free_mb >= vram_floor
|
|
# The RAM fallback is only meaningful where the device's memory *is*
|
|
# system memory. A discrete card falling back to RAM would mean running
|
|
# on the processor under the GPU's name.
|
|
ram_ok = device.uses_unified_memory and (ram_floor == 0 or device.ram_free_mb >= ram_floor)
|
|
if not (vram_ok or ram_ok):
|
|
return None
|
|
|
|
# Which resource carries it. A device with no dedicated VRAM at all that
|
|
# qualifies is running from system memory, and the placement says so,
|
|
# because "Mini-M4 · CPU · est. 47 min" and "oldbox · 3060 · est. 4 min"
|
|
# are the same job on very different hardware and a person needs to see
|
|
# which one they are waiting for.
|
|
via = "vram" if vram_ok else "ram"
|
|
|
|
would_evict: list[str] = []
|
|
for slug in requirement.models:
|
|
model = models[slug]
|
|
minimum = int(model["vram_min_mb"] or 0)
|
|
here = resident.get(slug)
|
|
if here is not None and here.get("device_id") in (device.id, str(device.id)):
|
|
continue # Already resident on this device. Free to use.
|
|
if minimum and vram_ok and device.vram_free_mb < minimum:
|
|
# Fits only if what is loaded here is unloaded first. Named rather
|
|
# than assumed: this is somebody's warm model.
|
|
would_evict.append(slug)
|
|
|
|
return Verdict(
|
|
via=via,
|
|
would_evict=tuple(sorted(would_evict)),
|
|
resident=not would_evict and via == "vram" and bool(requirement.models),
|
|
)
|
|
|
|
# -- the queue ---------------------------------------------------------
|
|
|
|
def submit_job(
|
|
self,
|
|
organization_id: UUID,
|
|
*,
|
|
kind: str,
|
|
requested_by: str,
|
|
requires: ModelRequirement | None = None,
|
|
payload: Mapping[str, Any] | None = None,
|
|
priority: int = 100,
|
|
project_id: UUID | None = None,
|
|
lang: str | None = None,
|
|
idempotency_key: str | None = None,
|
|
parent_job_id: UUID | None = None,
|
|
max_attempts: int = 3,
|
|
) -> Job:
|
|
"""Add work, or return the existing job for an idempotency key.
|
|
|
|
Re-running an import or a submission must not produce two rows, so the
|
|
key is a partial unique index rather than a convention.
|
|
"""
|
|
|
|
requirement = requires or ModelRequirement()
|
|
with self._connection() as connection, self._tx(connection):
|
|
if idempotency_key:
|
|
existing = connection.execute(
|
|
"SELECT * FROM public.jobs WHERE idempotency_key = %s", (idempotency_key,)
|
|
).fetchone()
|
|
if existing is not None:
|
|
return _job_from_row(existing)
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.jobs
|
|
(organization_id, project_id, lang, kind, status, requested_by,
|
|
payload, requires, priority, idempotency_key, parent_job_id, max_attempts)
|
|
VALUES (%s, %s, %s, %s, 'queued', %s, %s::jsonb, %s::jsonb, %s, %s, %s, %s)
|
|
RETURNING *
|
|
""",
|
|
(
|
|
organization_id,
|
|
project_id,
|
|
lang,
|
|
kind,
|
|
requested_by,
|
|
json.dumps(dict(payload or {})),
|
|
json.dumps(requirement.to_json()),
|
|
priority,
|
|
idempotency_key,
|
|
parent_job_id,
|
|
max_attempts,
|
|
),
|
|
).fetchone()
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, event, detail) VALUES (%s, 'queued', %s::jsonb)",
|
|
(row["id"], json.dumps({"kind": kind, "requires": requirement.to_json()})),
|
|
)
|
|
return _job_from_row(row)
|
|
|
|
def claim_job(
|
|
self,
|
|
node_id: UUID,
|
|
organization_id: UUID,
|
|
*,
|
|
lease_seconds: int = DEFAULT_LEASE_SECONDS,
|
|
kinds: Sequence[str] | None = None,
|
|
) -> ClaimedJob | None:
|
|
"""Atomically take one job for a node.
|
|
|
|
``FOR UPDATE SKIP LOCKED`` is the load-bearing line. It lets several
|
|
workers claim several *different* rows with no coordination between
|
|
them; an ordinary ``SELECT`` would have two of them pick the same row and
|
|
one overwrite the other's claim. (I25)
|
|
|
|
Eligibility is applied inside the same statement rather than by fetching
|
|
candidates and choosing in Python, because between choosing and claiming
|
|
another node may have taken the row - and a job placed on a node that
|
|
cannot run it fails later, further from the cause.
|
|
"""
|
|
|
|
now = _now()
|
|
expires = now + timedelta(seconds=lease_seconds)
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"""
|
|
UPDATE public.jobs SET
|
|
status = 'running',
|
|
lease_owner = %s,
|
|
lease_expires_at = %s,
|
|
attempt = attempt + 1,
|
|
started_at = COALESCE(started_at, %s)
|
|
WHERE id IN (
|
|
SELECT j.id
|
|
FROM public.jobs j
|
|
WHERE j.status = 'queued'
|
|
AND j.not_before <= %s
|
|
AND j.organization_id = %s
|
|
AND (%s::text[] IS NULL OR j.kind = ANY(%s::text[]))
|
|
ORDER BY j.priority, j.not_before, j.created_at
|
|
FOR UPDATE SKIP LOCKED
|
|
LIMIT 1
|
|
)
|
|
RETURNING *
|
|
""",
|
|
(node_id, expires, now, now, organization_id, kinds, kinds),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, node_id, event, detail) "
|
|
"VALUES (%s, %s, 'claimed', %s::jsonb)",
|
|
(row["id"], node_id, json.dumps({"attempt": row["attempt"],
|
|
"lease_expires_at": expires.isoformat()})),
|
|
)
|
|
job = _job_from_row(row)
|
|
# The placement is advisory here: the claim already happened, and a node
|
|
# that cannot run the job reports the failure rather than the scheduler
|
|
# pretending it placed work well.
|
|
placement = self.place(job.requires, organization_id=organization_id, prefer_node=node_id)
|
|
if placement is None or placement.node.id != node_id:
|
|
fallback = Placement(
|
|
node=ComputeNode(
|
|
id=node_id,
|
|
organization_id=organization_id,
|
|
name="(claiming node)",
|
|
is_controller=False,
|
|
trust_domain="local",
|
|
status="online",
|
|
),
|
|
device=Device(
|
|
id=uuid4(),
|
|
compute_node_id=node_id,
|
|
kind="cpu",
|
|
label="cpu:0",
|
|
vram_total_mb=0,
|
|
vram_free_mb=0,
|
|
ram_total_mb=0,
|
|
ram_free_mb=0,
|
|
),
|
|
reason="claimed_without_placement",
|
|
)
|
|
placement = fallback
|
|
return ClaimedJob(job=job, placement=placement, lease_expires_at=expires)
|
|
|
|
def renew_lease(self, job_id: UUID, node_id: UUID, *, lease_seconds: int = DEFAULT_LEASE_SECONDS) -> datetime | None:
|
|
"""Extend a lease while the work continues.
|
|
|
|
A lease is a liveness signal, not a timeout for the work. A long render
|
|
renews every few seconds; the lease expires only when the node actually
|
|
stops renewing, which is when the job is genuinely lost.
|
|
"""
|
|
|
|
expires = _now() + timedelta(seconds=lease_seconds)
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"UPDATE public.jobs SET lease_expires_at = %s "
|
|
"WHERE id = %s AND lease_owner = %s AND status = 'running' RETURNING lease_expires_at",
|
|
(expires, job_id, node_id),
|
|
).fetchone()
|
|
return row["lease_expires_at"] if row else None
|
|
|
|
def complete_job(self, job_id: UUID, node_id: UUID, *, result: Mapping[str, Any] | None = None) -> Job:
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"""
|
|
UPDATE public.jobs SET status = 'complete', result = %s::jsonb,
|
|
finished_at = now(), lease_owner = NULL, lease_expires_at = NULL,
|
|
error = NULL
|
|
WHERE id = %s AND lease_owner = %s AND status = 'running'
|
|
RETURNING *
|
|
""",
|
|
(json.dumps(dict(result or {})), job_id, node_id),
|
|
).fetchone()
|
|
if row is None:
|
|
raise ConflictError(
|
|
"job_not_claimed_by_this_node",
|
|
"The job is not held by this node, so it cannot be completed. It may "
|
|
"have been requeued after a lease expiry.",
|
|
)
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, node_id, event, detail) "
|
|
"VALUES (%s, %s, 'completed', %s::jsonb)",
|
|
(job_id, node_id, json.dumps(dict(result or {}))),
|
|
)
|
|
job = _job_from_row(row)
|
|
self._settle_parent(connection_organization_id=None, job=job)
|
|
return job
|
|
|
|
def fail_job(
|
|
self, job_id: UUID, node_id: UUID, *, error: str, retry: bool = True
|
|
) -> Job:
|
|
"""Record a failure, and decide whether the job gets another attempt.
|
|
|
|
A job that exhausts ``max_attempts`` ends ``failed`` with its last error
|
|
rather than retrying forever. A released lease is not a failure - see
|
|
:meth:`sweep_expired_leases` - so this path is for a job the node tried
|
|
and could not do.
|
|
"""
|
|
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"SELECT * FROM public.jobs WHERE id = %s AND lease_owner = %s",
|
|
(job_id, node_id),
|
|
).fetchone()
|
|
if row is None:
|
|
raise ConflictError(
|
|
"job_not_claimed_by_this_node",
|
|
"The job is not held by this node.",
|
|
)
|
|
if retry and row["attempt"] < row["max_attempts"]:
|
|
delay = min(BACKOFF_BASE_SECONDS * (2 ** max(row["attempt"] - 1, 0)), BACKOFF_CEILING_SECONDS)
|
|
updated = connection.execute(
|
|
"""
|
|
UPDATE public.jobs SET status = 'queued', lease_owner = NULL,
|
|
lease_expires_at = NULL, not_before = now() + make_interval(secs => %s),
|
|
error = %s
|
|
WHERE id = %s RETURNING *
|
|
""",
|
|
(delay, error, job_id),
|
|
).fetchone()
|
|
event = "requeued"
|
|
else:
|
|
updated = connection.execute(
|
|
"""
|
|
UPDATE public.jobs SET status = 'failed', lease_owner = NULL,
|
|
lease_expires_at = NULL, finished_at = now(), error = %s
|
|
WHERE id = %s RETURNING *
|
|
""",
|
|
(error, job_id),
|
|
).fetchone()
|
|
event = "failed"
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, node_id, event, detail) "
|
|
"VALUES (%s, %s, %s, %s::jsonb)",
|
|
(job_id, node_id, event, json.dumps({"error": error})),
|
|
)
|
|
return _job_from_row(updated)
|
|
|
|
def submit_children(
|
|
self, parent_job_id: UUID, node_id: UUID, children: Sequence[Mapping[str, Any]]
|
|
) -> list[Job]:
|
|
"""Decompose a job the node decided was too large (§6.19, I29).
|
|
|
|
Children are ordinary jobs with a ``parent_job_id``, so the scheduler
|
|
needs no special case and the queue needs no second table. A container is
|
|
never itself dispatched.
|
|
"""
|
|
|
|
created: list[Job] = []
|
|
with self._connection() as connection, self._tx(connection):
|
|
parent = connection.execute(
|
|
"SELECT * FROM public.jobs WHERE id = %s", (parent_job_id,)
|
|
).fetchone()
|
|
if parent is None:
|
|
raise NotFoundError("job_not_found")
|
|
if parent["lease_owner"] != node_id:
|
|
raise ConflictError(
|
|
"job_not_claimed_by_this_node", "Only the holding node may decompose a job."
|
|
)
|
|
for spec in children:
|
|
requirement = ModelRequirement.from_json(spec.get("requires"))
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.jobs
|
|
(organization_id, project_id, lang, kind, status, requested_by,
|
|
payload, requires, priority, parent_job_id, idempotency_key)
|
|
VALUES (%s, %s, %s, %s, 'queued', %s, %s::jsonb, %s::jsonb, %s, %s, %s)
|
|
RETURNING *
|
|
""",
|
|
(
|
|
parent["organization_id"],
|
|
parent["project_id"],
|
|
parent["lang"],
|
|
spec["kind"],
|
|
f"node:{node_id}",
|
|
json.dumps(dict(spec.get("payload", {}))),
|
|
json.dumps(requirement.to_json()),
|
|
int(spec.get("priority", parent["priority"])),
|
|
parent_job_id,
|
|
spec.get("idempotency_key"),
|
|
),
|
|
).fetchone()
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, event, detail) "
|
|
"VALUES (%s, 'decomposed', %s::jsonb)",
|
|
(row["id"], json.dumps({"parent": str(parent_job_id)})),
|
|
)
|
|
created.append(_job_from_row(row))
|
|
# The parent becomes a container: released, and not claimable.
|
|
connection.execute(
|
|
"UPDATE public.jobs SET status = 'running', lease_owner = NULL, "
|
|
"lease_expires_at = NULL WHERE id = %s",
|
|
(parent_job_id,),
|
|
)
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, node_id, event, detail) "
|
|
"VALUES (%s, %s, 'decomposed', %s::jsonb)",
|
|
(parent_job_id, node_id, json.dumps({"children": len(created)})),
|
|
)
|
|
return created
|
|
|
|
def sweep_expired_leases(self, *, now: datetime | None = None) -> list[UUID]:
|
|
"""Return jobs whose holder stopped renewing to the queue.
|
|
|
|
This is the recovery path, and it exists because a worker that is
|
|
``kill -9``'d never releases anything. The lease is what releases it.
|
|
|
|
The status becomes ``queued`` rather than ``failed``, because nothing
|
|
about an expired lease says the work was impossible - only that whoever
|
|
had it is gone. ``attempt`` has already been incremented at claim time,
|
|
so a job that kills a worker repeatedly still terminates. (I26)
|
|
"""
|
|
|
|
moment = now or _now()
|
|
with self._connection() as connection, self._tx(connection):
|
|
rows = connection.execute(
|
|
"""
|
|
UPDATE public.jobs SET
|
|
status = CASE WHEN attempt >= max_attempts THEN 'failed' ELSE 'queued' END,
|
|
lease_owner = NULL,
|
|
lease_expires_at = NULL,
|
|
not_before = now() + make_interval(secs => %s),
|
|
error = 'lease expired: the holding node stopped renewing',
|
|
finished_at = CASE WHEN attempt >= max_attempts THEN now() ELSE NULL END
|
|
WHERE status = 'running'
|
|
AND lease_expires_at IS NOT NULL
|
|
AND lease_expires_at < %s
|
|
RETURNING id, status, attempt
|
|
""",
|
|
(BACKOFF_BASE_SECONDS, moment),
|
|
).fetchall()
|
|
for row in rows:
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, event, detail) "
|
|
"VALUES (%s, 'lease_expired', %s::jsonb)",
|
|
(row["id"], json.dumps({"attempt": row["attempt"], "now": row["status"]})),
|
|
)
|
|
return [row["id"] for row in rows]
|
|
|
|
def mark_blocked(self, job_id: UUID, reason: str) -> Job:
|
|
"""Record that nothing can run this job, and say why. (§7.3.1)
|
|
|
|
A job that cannot run must not sit in the queue looking like it is
|
|
progressing. The difference between "slow" and "impossible" is the
|
|
difference between a long wait and a silent hang, and today
|
|
``gemma4:26b`` is only-runnable on oldbox, so this is not hypothetical.
|
|
"""
|
|
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"UPDATE public.jobs SET status = 'blocked', error = %s, finished_at = now() "
|
|
"WHERE id = %s AND status = 'queued' RETURNING *",
|
|
(reason, job_id),
|
|
).fetchone()
|
|
if row is None:
|
|
raise ConflictError("job_not_blockable", "Only a queued job can be blocked.")
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, event, detail) VALUES (%s, 'blocked', %s::jsonb)",
|
|
(job_id, json.dumps({"reason": reason})),
|
|
)
|
|
return _job_from_row(row)
|
|
|
|
def get_job(self, job_id: UUID) -> Job | None:
|
|
with self._connection() as connection:
|
|
row = connection.execute("SELECT * FROM public.jobs WHERE id = %s", (job_id,)).fetchone()
|
|
return _job_from_row(row) if row else None
|
|
|
|
def list_jobs(self, organization_id: UUID, *, statuses: Sequence[str] | None = None) -> list[Job]:
|
|
query = "SELECT * FROM public.jobs WHERE organization_id = %s"
|
|
params: list[Any] = [organization_id]
|
|
if statuses:
|
|
query += " AND status = ANY(%s)"
|
|
params.append(list(statuses))
|
|
query += " ORDER BY created_at"
|
|
with self._connection() as connection:
|
|
rows = connection.execute(query, params).fetchall()
|
|
return [_job_from_row(row) for row in rows]
|
|
|
|
def job_events(self, job_id: UUID) -> list[dict[str, Any]]:
|
|
with self._connection() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM public.job_events WHERE job_id = %s ORDER BY at, id", (job_id,)
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
def _settle_parent(self, *, connection_organization_id: UUID | None, job: Job) -> None:
|
|
"""Roll a child's outcome up to its container. (I29)
|
|
|
|
A parent with any failed child ends ``partial``, never ``complete``:
|
|
"3 of 49 chapters analysed" and "failed" are different facts and
|
|
collapsing them loses the one that matters.
|
|
"""
|
|
|
|
if job.parent_job_id is None:
|
|
return
|
|
with self._connection() as connection, self._tx(connection):
|
|
children = connection.execute(
|
|
"SELECT status, count(*) AS total FROM public.jobs WHERE parent_job_id = %s GROUP BY status",
|
|
(job.parent_job_id,),
|
|
).fetchall()
|
|
if not children:
|
|
return
|
|
counts = {row["status"]: row["total"] for row in children}
|
|
if any(status not in TERMINAL_STATUSES for status in counts):
|
|
return
|
|
failed = counts.get(FAILED, 0) + counts.get(CANCELLED, 0)
|
|
if failed:
|
|
outcome, detail = PARTIAL, f"{failed} of {sum(counts.values())} children failed"
|
|
elif counts.get(PARTIAL):
|
|
outcome, detail = PARTIAL, "one or more children ended partial"
|
|
else:
|
|
outcome, detail = COMPLETE, "all children complete"
|
|
connection.execute(
|
|
"UPDATE public.jobs SET status = %s, finished_at = now(), error = %s WHERE id = %s",
|
|
(outcome, None if outcome == COMPLETE else detail, job.parent_job_id),
|
|
)
|
|
connection.execute(
|
|
"INSERT INTO public.job_events (job_id, event, detail) VALUES (%s, 'settled', %s::jsonb)",
|
|
(job.parent_job_id, json.dumps({"outcome": outcome, "detail": detail})),
|
|
)
|
|
|
|
def close_database(self) -> None: # pragma: no cover - convenience alias
|
|
self.close()
|
|
|
|
|
|
# -- row mapping -------------------------------------------------------------
|
|
|
|
|
|
def now_aware(value: datetime | None) -> datetime:
|
|
"""Coerce a timestamp to timezone-aware UTC.
|
|
|
|
``timestamptz`` comes back aware, but a value that has crossed a
|
|
serialisation boundary may not be, and comparing the two raises rather than
|
|
returning a wrong answer. Coercing is the safe direction: it can only make
|
|
the staleness check stricter, never looser.
|
|
"""
|
|
|
|
if value is None:
|
|
return _now()
|
|
return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
|
|
|
|
|
|
def _device_from_row(row: Mapping[str, Any]) -> Device:
|
|
return Device(
|
|
id=row["id"],
|
|
compute_node_id=row["compute_node_id"],
|
|
kind=row["kind"],
|
|
label=row["label"],
|
|
vram_total_mb=int(row["vram_total_mb"] or 0),
|
|
vram_free_mb=int(row["vram_free_mb"] or 0),
|
|
ram_total_mb=int(row["ram_total_mb"] or 0),
|
|
ram_free_mb=int(row["ram_free_mb"] or 0),
|
|
compute_units=float(row["compute_units"]) if row.get("compute_units") is not None else None,
|
|
throughput_factor=float(row.get("throughput_factor") or 1.0),
|
|
supports_accelerators=tuple(row.get("supports_accelerators") or ()),
|
|
)
|
|
|
|
|
|
def _node_from_row(row: Mapping[str, Any], *, devices: Sequence[Device]) -> ComputeNode:
|
|
return ComputeNode(
|
|
id=row["id"],
|
|
organization_id=row["organization_id"],
|
|
name=row["name"],
|
|
is_controller=bool(row.get("is_controller")),
|
|
trust_domain=row.get("trust_domain") or "local",
|
|
status=row["status"],
|
|
locality=row.get("locality"),
|
|
transport=row.get("transport") or "https",
|
|
agent_version=row.get("agent_version"),
|
|
agent_started_at=row.get("agent_started_at"),
|
|
last_heartbeat_at=row.get("last_heartbeat_at"),
|
|
created_at=row.get("created_at"),
|
|
devices=tuple(devices),
|
|
)
|
|
|
|
|
|
def _job_from_row(row: Mapping[str, Any]) -> Job:
|
|
return Job(
|
|
id=row["id"],
|
|
organization_id=row["organization_id"],
|
|
kind=row["kind"],
|
|
status=row["status"],
|
|
requested_by=row["requested_by"],
|
|
priority=int(row.get("priority") or 100),
|
|
requires=ModelRequirement.from_json(row.get("requires")),
|
|
payload=row.get("payload") or {},
|
|
project_id=row.get("project_id"),
|
|
lang=row.get("lang"),
|
|
lease_owner=row.get("lease_owner"),
|
|
lease_expires_at=row.get("lease_expires_at"),
|
|
attempt=int(row.get("attempt") or 0),
|
|
max_attempts=int(row.get("max_attempts") or 3),
|
|
not_before=row.get("not_before"),
|
|
parent_job_id=row.get("parent_job_id"),
|
|
idempotency_key=row.get("idempotency_key"),
|
|
result=row.get("result"),
|
|
error=row.get("error"),
|
|
created_at=row.get("created_at"),
|
|
started_at=row.get("started_at"),
|
|
finished_at=row.get("finished_at"),
|
|
)
|
|
|
|
|
|
def _placement_reason(
|
|
device: Device, prefer_node: UUID | None, is_preferred: bool, slow_path: bool
|
|
) -> str:
|
|
"""A short, machine-readable account of why this placement was chosen.
|
|
|
|
Recorded on the job rather than implied, because a scheduler that silently
|
|
prefers one machine cannot be debugged when it starves another. A person
|
|
looking at a queue needs to see that a job is on the slow path and roughly
|
|
what that costs.
|
|
"""
|
|
|
|
prefix = "local_affinity" if (is_preferred and prefer_node is not None) else "placed"
|
|
if slow_path:
|
|
return f"{prefix}:slow_path:ram_not_vram"
|
|
if device.kind == "cpu":
|
|
return f"{prefix}:cpu"
|
|
return f"{prefix}:{device.kind}"
|
|
|
|
|
|
def _placement_score(placement: Placement) -> float:
|
|
"""Rank placements. Higher wins.
|
|
|
|
Throughput dominates, residency discounts, and a placement that would evict
|
|
somebody's model is discounted hardest. The reason a warm model beats a
|
|
marginally faster cold one is that loading is not free: minutes of held
|
|
VRAM before the first token, and a queue behind the load.
|
|
"""
|
|
|
|
score = placement.device.throughput_factor
|
|
if placement.would_evict:
|
|
score *= 0.4
|
|
elif not placement.resident:
|
|
score *= 0.7
|
|
if placement.reason.startswith("local_affinity"):
|
|
score *= 1.5
|
|
return score
|
|
|
|
|
|
__all__ = [
|
|
"BACKOFF_BASE_SECONDS",
|
|
"BLOCKED",
|
|
"COMPLETE",
|
|
"ClaimedJob",
|
|
"ComputeNode",
|
|
"ComputeRepository",
|
|
"DEFAULT_LEASE_SECONDS",
|
|
"Device",
|
|
"FAILED",
|
|
"HEARTBEAT_STALE_SECONDS",
|
|
"Job",
|
|
"ModelRequirement",
|
|
"NODE_SCOPES",
|
|
"PARTIAL",
|
|
"Placement",
|
|
"QUEUED",
|
|
"RUNNING",
|
|
"TERMINAL_STATUSES",
|
|
]
|
|
|
|
|
|
def build_compute_repository(settings, passwords=None) -> "ComputeRepository":
|
|
"""Build the compute repository from settings.
|
|
|
|
Shared by the request dependency and by the application's startup, because
|
|
the lease sweeper needs a repository before any request arrives. Building it
|
|
lazily in the sweeper instead would mean a controller restarted mid-job never
|
|
recovers that job: nothing would ever touch the compute surface, so the
|
|
repository would never be built, so the sweeper would have nothing to sweep,
|
|
so the job would sit ``running`` for ever - which is the failure the sweeper
|
|
exists to prevent, reached by the route the sweeper takes to prevent it.
|
|
"""
|
|
|
|
if not getattr(settings, "database_dsn", None):
|
|
raise RepositoryError("compute_unavailable", "No database is configured.")
|
|
return ComputeRepository(settings.database_dsn, passwords=passwords)
|