Replaces claiming with assignment. The controller hands a job to one or more nodes and a node asks what it has been given; jobs.lease_owner could name only one holder, so a job fanned out to several machines was inexpressible, and 'what was assigned and what is not complete' had no answer. job_assignments is one row per (job, node, attempt), so both are now readable. Nodes poll rather than being pushed, which is not a simplification: nodes sit behind NAT and the controller has no route to them. The only path that exists is a node calling out. The rule with teeth: a job is complete when the controller has the artifacts to show it. Reporting is not believing. A node that finishes and hands over nothing gets its assignment failed, with the reason naming what was missing, and the work stays available for another machine. The tempting implementation - mark complete, then verify afterwards - cannot fail, because the only thing it could compare against is the report it just trusted. artifacts_verified_at is written by verification alone, so 'the node says it is done' and 'the controller confirmed it' are different facts that cannot look alike. fan_out defaults to 1 and is never inferred. More than one is redundant execution: several machines do the same work and the first verified result wins, the rest are cancelled. Worth having for cheap idempotent work, wasteful for anything costing GPU-seconds, so a caller has to say so. The losers are cancelled rather than left running - otherwise every machine on the list burns its time on a result nobody will read. Four defects found while building it, all by tests written first: - Naming a node explicitly bypassed the drain check. An administrator could drain a machine, see it read as drained, and still hand it work: the control doing nothing while appearing to work. Status is now checked on both routes. - attempts were counted from jobs.attempt, which only claim_job incremented, so under assignments max_attempts was never reached and a failing job was re-offered forever. Counted from failed assignments now, so a fan-out to three costs one attempt, not three. - start_assignment reported 'not startable' when the real problem was that the assignment belonged to another node. Ownership is checked first, because an operator handed the first will look in the wrong place. - The artifact API accepted values the database rejects: role 'output' where job_artifacts_role_check allows only 'consumed' and 'produced', an unbounded status, a content_hash that was not 64 hex characters, and kind at 80 where the limit is 64. Every completion would have been a 500. The schema had fixed this vocabulary all along and the API had to learn it rather than invent it. Artifacts are recorded even when verification fails: a short run still produced evidence, and discarding it would destroy the record of why. storage_key is never treated as proof the bytes exist, because there is no object storage to ask yet (DECISION-6). Claim still exists. It is listed in the pinned node surface as pending removal and the next change takes it out - two paths handing out work is the state this project keeps finding and criticising. Mutation-verified: trusting the node fails 3 tests, dropping the drain check fails 1, leaving fan-out losers running fails 1. 674 passed, 4 skipped.
2975 lines
126 KiB
Python
2975 lines
126 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
|
|
import secrets
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any, Callable, Mapping, Sequence
|
|
from uuid import UUID, uuid4
|
|
|
|
import psycopg
|
|
from psycopg import sql
|
|
|
|
from .errors import (
|
|
ConflictError,
|
|
InputError,
|
|
NodeNameExhaustedError,
|
|
NotFoundError,
|
|
RepositoryError,
|
|
)
|
|
from .security import PasswordService, generate_token
|
|
from .tenancy import tenant_schema_name
|
|
|
|
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
|
|
|
|
#: Pairing guesses allowed per source address per 15-minute window. Twenty is enough
|
|
#: for a person who mistypes a six-digit code a few times and needs no more, and small
|
|
#: enough that the 1,000,000-code space cannot be walked inside one window: at twenty
|
|
#: per quarter-hour the space needs about 21 years. The window resets rather than
|
|
#: decaying, so this is a rate and not a total. Without it a six-digit code is not
|
|
#: defensible, however well the rest of the flow is built.
|
|
PAIRING_ATTEMPT_BUDGET = 20
|
|
|
|
#: 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 JobAssignment:
|
|
"""One machine's attempt at one job.
|
|
|
|
Distinct from the job on purpose. A job fanned out to three machines has three of
|
|
these and one ``Job``, and the difference between "the job is done" and "this
|
|
machine thinks it is done" is exactly what ``artifacts_verified_at`` records.
|
|
"""
|
|
|
|
id: UUID
|
|
job_id: UUID
|
|
node_id: UUID
|
|
state: str
|
|
attempt: int = 1
|
|
assigned_at: datetime | None = None
|
|
started_at: datetime | None = None
|
|
finished_at: datetime | None = None
|
|
lease_expires_at: datetime | None = None
|
|
artifacts_verified_at: datetime | None = None
|
|
result: Mapping[str, Any] = field(default_factory=dict)
|
|
error: str | None = None
|
|
|
|
@property
|
|
def is_verified(self) -> bool:
|
|
"""Whether the controller has confirmed this assignment with artifacts."""
|
|
|
|
return self.artifacts_verified_at is not None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class AssignmentOutcome:
|
|
"""What came back from an assignment operation, including the part that failed.
|
|
|
|
``missing`` is the list of things verification wanted and did not find. It is in the
|
|
return value rather than only in the error message because the caller - a node, or an
|
|
administrator watching a console - needs to know *what* was absent to act on it, and
|
|
parsing that back out of a sentence is how it gets ignored.
|
|
"""
|
|
|
|
assignment: JobAssignment
|
|
verified: bool
|
|
job_status: str
|
|
missing: Sequence[str] = ()
|
|
artifacts_recorded: Sequence[str] = ()
|
|
superseded_nodes: Sequence[str] = ()
|
|
|
|
|
|
def _assignment_from_row(row: Mapping[str, Any]) -> JobAssignment:
|
|
return JobAssignment(
|
|
id=row["id"],
|
|
job_id=row["job_id"],
|
|
node_id=row["node_id"],
|
|
state=row["state"],
|
|
attempt=int(row["attempt"]),
|
|
assigned_at=row.get("assigned_at"),
|
|
started_at=row.get("started_at"),
|
|
finished_at=row.get("finished_at"),
|
|
lease_expires_at=row.get("lease_expires_at"),
|
|
artifacts_verified_at=row.get("artifacts_verified_at"),
|
|
result=dict(row.get("result") or {}),
|
|
error=row.get("error"),
|
|
)
|
|
|
|
|
|
#: The letters used to separate two machines that share a name *and* the last four hex
|
|
#: digits of their MAC. Twenty-six is enough for any realistic collision and the
|
|
#: sequence continues past Z rather than stopping, because a limit that silently runs
|
|
#: out is worse than a long name: see :func:`_letter_suffix`.
|
|
SUFFIX_LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
|
|
|
|
|
|
def _letter_suffix(index: int) -> str:
|
|
"""``0 -> A``, ``25 -> Z``, ``26 -> AA``.
|
|
|
|
Spreadsheet-style, so the sequence does not run out at 26 and quietly start reusing
|
|
letters. A collision that cannot be resolved should be reported as a collision, not
|
|
dressed up as a fresh name.
|
|
"""
|
|
|
|
if index < 0: # pragma: no cover - callers pass a counter
|
|
raise ValueError("suffix index cannot be negative")
|
|
letters = ""
|
|
index += 1
|
|
while index > 0:
|
|
index, remainder = divmod(index - 1, len(SUFFIX_LETTERS))
|
|
letters = SUFFIX_LETTERS[remainder] + letters
|
|
return letters
|
|
|
|
|
|
def mac_suffix(mac: str | None) -> str | None:
|
|
"""The last four hex digits of a MAC, uppercased. ``None`` if there isn't one.
|
|
|
|
``...:ac:4b`` becomes ``AC4B``, so ``VM238`` on that machine is ``VM238-AC4B``.
|
|
Four digits rather than the whole address because a node's name appears in a UI
|
|
column, a command line and a systemd unit, and the whole thing would be unreadable
|
|
in all three. Note it is the last four *characters of the address*, so it is the
|
|
final two octets - not four octets, and not one octet padded out.
|
|
|
|
Sixteen bits of entropy is not much - 65 536 machines could share a base name and a
|
|
suffix - so this narrows collisions rather than eliminating them, and
|
|
:func:`resolve_node_name` finishes the job.
|
|
"""
|
|
|
|
if mac is None:
|
|
return None
|
|
# Six octets of two hex digits. Validated rather than cleaned up, because the
|
|
# permissive version of this - strip the separators and take the last four characters -
|
|
# turns "not-a-mac" into AMAC, and a node named VM238-AMAC is worse than a node that
|
|
# admits it does not know its own address.
|
|
cleaned = mac.strip().replace("-", ":").lower()
|
|
parts = cleaned.split(":")
|
|
if len(parts) != 6 or not all(len(part) == 2 for part in parts):
|
|
return None
|
|
if not all(character in "0123456789abcdef" for part in parts for character in part):
|
|
return None
|
|
return f"{parts[4]}{parts[5]}".upper()
|
|
|
|
|
|
def resolve_node_name(
|
|
base_name: str,
|
|
mac: str | None,
|
|
taken: Callable[[str], bool],
|
|
) -> str:
|
|
"""The name this machine gets: ``base-AC4B``, or ``base-AC4B-B`` and so on.
|
|
|
|
``base_name`` is what the operator typed, and is what the pairing code was minted
|
|
against. The suffix is what makes two machines called ``VM238`` distinguishable, and
|
|
it comes from hardware the controller cannot know in advance - which is why the code
|
|
binds to the base name and the final name is resolved here, at redeem, where the
|
|
uniqueness check can be made against the database rather than guessed at by the
|
|
script running on the machine.
|
|
|
|
The letter fallback exists because four hex digits is sixteen bits and two machines
|
|
can genuinely share it: cloned VM images with the same generated MAC, or a
|
|
hypervisor that hands out the same address twice. Without it, the second machine to
|
|
pair would be refused for a name that looks free to the person typing it.
|
|
|
|
``taken`` is called with a candidate and must answer whether that name already
|
|
exists. It is injected rather than queried here so this stays a pure function: the
|
|
rule can then be tested exhaustively without a database, and the check-and-create
|
|
race is closed at the insert by the unique constraint rather than here.
|
|
"""
|
|
|
|
suffix = mac_suffix(mac)
|
|
stem = f"{base_name}-{suffix}" if suffix else base_name
|
|
if not taken(stem):
|
|
return stem
|
|
for index in range(len(SUFFIX_LETTERS) * 4):
|
|
candidate = f"{stem}-{_letter_suffix(index)}"
|
|
if not taken(candidate):
|
|
return candidate
|
|
raise NodeNameExhaustedError(
|
|
"node_name_exhausted",
|
|
f"Every name derived from {base_name!r} is already in use. The base name is "
|
|
"too widely shared; pair with a more specific one.",
|
|
)
|
|
|
|
|
|
@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
|
|
|
|
#: What the node is *right now*, derived from the stored status and the clock by
|
|
#: `effective_status`. This is the value every caller should branch on: it is
|
|
#: `offline` the moment a node goes quiet, with no background task involved.
|
|
status: str
|
|
|
|
#: What the database column says, kept so an operator can see that the two differ
|
|
#: rather than wondering why a row reads `online` for a machine that is gone.
|
|
stored_status: str | None = None
|
|
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):
|
|
# Refused here rather than left to the database trigger, so the caller gets a
|
|
# reason it can read instead of a constraint violation. The trigger stays as
|
|
# the guarantee; this is the message.
|
|
existing = self._node_row(connection, node_id)
|
|
if existing is not None and existing["status"] == "removed":
|
|
raise RepositoryError(
|
|
"node_removed",
|
|
"This node has been removed from the fleet and cannot return to "
|
|
"service. A replacement machine pairs under its own name.",
|
|
)
|
|
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
|
|
|
|
def revoke_node_credentials(
|
|
self,
|
|
node_id: UUID,
|
|
*,
|
|
reason: str | None = None,
|
|
actor: str | None = None,
|
|
) -> int:
|
|
"""Revoke every credential this node holds. Returns how many were live.
|
|
|
|
The `revoked_at` column and the check in :meth:`resolve_node_credential` have been
|
|
there since pairing; what was missing was any way to set it. So the ability to
|
|
cut a node off existed on paper only, and the routes for removing one did not
|
|
exist at all.
|
|
|
|
Revoking rather than deleting, because the credential row is the record that a
|
|
node was ever onboarded and which key was issued. Deleting it would make a
|
|
compromised key indistinguishable from one that never existed.
|
|
|
|
Reason is recorded in the returned detail rather than logged, so the audit trail
|
|
says why a machine went dark instead of only that it did.
|
|
"""
|
|
|
|
with self._connection() as connection, self._tx(connection):
|
|
revoked = connection.execute(
|
|
"""
|
|
UPDATE public.node_credentials
|
|
SET revoked_at = now()
|
|
WHERE compute_node_id = %s AND revoked_at IS NULL
|
|
""",
|
|
(node_id,),
|
|
).rowcount
|
|
if self._node_row(connection, node_id) is None:
|
|
raise NotFoundError("compute_node_not_found")
|
|
if revoked:
|
|
self._node_event(
|
|
connection,
|
|
node_id,
|
|
"credentials_revoked",
|
|
{
|
|
"count": revoked,
|
|
"reason": reason,
|
|
"actor": actor,
|
|
},
|
|
)
|
|
return int(revoked or 0)
|
|
|
|
def remove_node(
|
|
self,
|
|
node_id: UUID,
|
|
*,
|
|
reason: str | None = None,
|
|
actor: str | None = None,
|
|
) -> ComputeNode:
|
|
"""Take a machine out of the fleet for good.
|
|
|
|
What this deletes is everything that could let the node *act*: its credentials,
|
|
its device measurements, its heartbeats and its model reports. What it keeps is
|
|
the `compute_nodes` row, and that is not tidiness.
|
|
|
|
The row cannot simply be deleted. `job_events.node_id` references it and is
|
|
written by every path that records work, and `node_pairing_codes.used_by_node_id`
|
|
references it too. Both are `ON DELETE NO ACTION`, so the delete is refused.
|
|
Changing them to `SET NULL` does not help either: that is an `UPDATE` of
|
|
`job_events`, and `job_events_append_only` raises on every update by design. So a
|
|
node that has ever done work cannot be removed without weakening the append-only
|
|
guarantee on the event trail - and that guarantee is worth more than a tidy
|
|
`DELETE`.
|
|
|
|
So the machine goes and the record stays. A `removed` node cannot authenticate,
|
|
cannot be assigned work, and is skipped by the claim path on status alone.
|
|
|
|
Work the node was holding is released here rather than left to the lease to
|
|
expire. An operator removing dead hardware knows the work did not finish; making
|
|
them wait out a lease to find that out is the wrong way round, and for a
|
|
long lease it is most of an hour of the queue stalled behind a machine nobody
|
|
is going to switch on.
|
|
"""
|
|
|
|
with self._connection() as connection, self._tx(connection):
|
|
if self._node_row(connection, node_id) is None:
|
|
raise NotFoundError("compute_node_not_found")
|
|
|
|
# Released before the credentials go, so the work is accounted for even if
|
|
# the node was never going to come back.
|
|
released = connection.execute(
|
|
"""
|
|
UPDATE public.jobs SET
|
|
status = CASE WHEN attempt >= max_attempts THEN 'failed' ELSE 'queued' END,
|
|
lease_owner = NULL,
|
|
lease_expires_at = NULL
|
|
WHERE lease_owner = %s AND status = 'running'
|
|
RETURNING id, status
|
|
""",
|
|
(node_id,),
|
|
).fetchall()
|
|
for row in released:
|
|
self._job_event(
|
|
connection,
|
|
row["id"],
|
|
"released_node_removed",
|
|
{"node_id": str(node_id), "status": row["status"], "reason": reason},
|
|
)
|
|
|
|
# CASCADE would do all four of these from the node row, but the node row is
|
|
# staying, so they are named explicitly.
|
|
for table in ("compute_devices", "node_heartbeats", "node_models"):
|
|
connection.execute(f"DELETE FROM public.{table} WHERE compute_node_id = %s", (node_id,))
|
|
connection.execute(
|
|
"UPDATE public.node_credentials SET revoked_at = now() "
|
|
"WHERE compute_node_id = %s AND revoked_at IS NULL",
|
|
(node_id,),
|
|
)
|
|
|
|
connection.execute(
|
|
"UPDATE public.compute_nodes SET status = 'removed', updated_at = now() "
|
|
"WHERE id = %s",
|
|
(node_id,),
|
|
)
|
|
self._node_event(
|
|
connection,
|
|
node_id,
|
|
"removed",
|
|
{
|
|
"reason": reason,
|
|
"actor": actor,
|
|
"released_jobs": [str(row["id"]) for row in released],
|
|
},
|
|
)
|
|
node = self.get_node(node_id)
|
|
assert node is not None
|
|
return node
|
|
|
|
# -- assignments -------------------------------------------------------
|
|
|
|
def assign_job(
|
|
self,
|
|
job_id: UUID,
|
|
*,
|
|
node_ids: Sequence[UUID] | None = None,
|
|
kinds: Sequence[str] | None = None,
|
|
fan_out: int = 1,
|
|
lease_seconds: int = DEFAULT_LEASE_SECONDS,
|
|
actor: str | None = None,
|
|
) -> list[JobAssignment]:
|
|
"""The controller hands a job to one or more nodes.
|
|
|
|
This replaces `claim_job`. Work now reaches a node because the controller put it
|
|
there, not because the node decided it could use the job - which is what makes
|
|
"the controller knows what was assigned and what is not complete" answerable, and
|
|
what makes a job fanning out to several machines expressible at all.
|
|
`jobs.lease_owner` could only ever name one holder.
|
|
|
|
`fan_out` defaults to 1. More than one is redundant execution: several machines
|
|
do the same work and the first whose artifacts verify wins, the rest are
|
|
cancelled. That is worth having for cheap idempotent work and expensive for
|
|
anything that costs GPU-seconds, so it is never assumed - a caller has to say so.
|
|
|
|
`node_ids` names the machines outright; `kinds` asks for whoever matches the job's
|
|
kind. Both may be given, in which case `node_ids` wins, because naming a machine
|
|
is a stronger statement than naming a kind.
|
|
|
|
The job's own status does not move. It stays `queued` until an assignment actually
|
|
starts, because a job nobody has picked up yet is not running work, and a console
|
|
that says otherwise is claiming progress that has not happened.
|
|
"""
|
|
|
|
if fan_out < 1:
|
|
raise InputError("fan_out must be at least 1")
|
|
|
|
now = _now()
|
|
expires = now + timedelta(seconds=lease_seconds)
|
|
with self._connection() as connection, self._tx(connection):
|
|
job = self._job_row(connection, job_id, for_update=True)
|
|
if job is None:
|
|
raise NotFoundError("job_not_found")
|
|
if job["status"] in TERMINAL_STATUSES:
|
|
raise ConflictError(
|
|
"job_already_finished",
|
|
f"The job is {job['status']!r} and cannot be assigned.",
|
|
)
|
|
|
|
candidates = self._candidate_nodes(
|
|
connection,
|
|
node_ids=node_ids,
|
|
kinds=list(kinds) if kinds else [job["kind"]],
|
|
)
|
|
if not candidates:
|
|
return []
|
|
|
|
chosen = candidates[:fan_out]
|
|
created: list[JobAssignment] = []
|
|
for node_id in chosen:
|
|
attempt = self._next_attempt(connection, job_id, node_id)
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.job_assignments (
|
|
job_id, node_id, state, attempt, assigned_at, lease_expires_at
|
|
)
|
|
VALUES (%s, %s, 'assigned', %s, %s, %s)
|
|
ON CONFLICT DO NOTHING
|
|
RETURNING *
|
|
""",
|
|
(job_id, node_id, attempt, now, expires),
|
|
).fetchone()
|
|
if row is None:
|
|
# Already holds an open assignment for this job. Not an error: the
|
|
# caller asked for it to be assigned and it is.
|
|
continue
|
|
self._job_event(
|
|
connection,
|
|
job_id,
|
|
"assigned",
|
|
{"node_id": str(node_id), "attempt": attempt},
|
|
node_id=node_id,
|
|
)
|
|
created.append(_assignment_from_row(row))
|
|
|
|
if created:
|
|
self._node_event(
|
|
connection,
|
|
chosen[0],
|
|
"job_assigned",
|
|
{"job_id": str(job_id), "count": len(created), "actor": actor},
|
|
)
|
|
return created
|
|
|
|
def assignments_for_node(
|
|
self,
|
|
node_id: UUID,
|
|
*,
|
|
limit: int = 10,
|
|
include_running: bool = False,
|
|
) -> list[JobAssignment]:
|
|
"""What this node has been asked to do. The node's whole view of the queue.
|
|
|
|
Returns assignments in priority order rather than assignment order, because the
|
|
controller expresses urgency through `jobs.priority` and a node that took them in
|
|
the order they were handed out would ignore it.
|
|
|
|
`include_running` is off by default: an agent polls in a loop, and returning the
|
|
assignment it is already working on would have it start the same work twice on
|
|
every poll.
|
|
"""
|
|
|
|
states = ("assigned", "running") if include_running else ("assigned",)
|
|
with self._connection() as connection:
|
|
rows = connection.execute(
|
|
"""
|
|
SELECT a.* FROM public.job_assignments a
|
|
JOIN public.jobs j ON j.id = a.job_id
|
|
WHERE a.node_id = %s
|
|
AND a.state = ANY(%s::text[])
|
|
AND (a.lease_expires_at IS NULL OR a.lease_expires_at > now())
|
|
ORDER BY j.priority, j.not_before, a.assigned_at
|
|
LIMIT %s
|
|
""",
|
|
(node_id, list(states), limit),
|
|
).fetchall()
|
|
return [_assignment_from_row(row) for row in rows]
|
|
|
|
def assignments_for_job(self, job_id: UUID) -> list[JobAssignment]:
|
|
"""Every attempt made on this job, oldest first.
|
|
|
|
"What was assigned, and what is not complete" is answered by this list and nothing
|
|
else: one row per attempt, each saying which machine, when, how it ended, and
|
|
whether the controller confirmed it with artifacts.
|
|
"""
|
|
|
|
with self._connection() as connection:
|
|
rows = connection.execute(
|
|
"SELECT * FROM public.job_assignments WHERE job_id = %s "
|
|
"ORDER BY assigned_at, attempt",
|
|
(job_id,),
|
|
).fetchall()
|
|
return [_assignment_from_row(row) for row in rows]
|
|
|
|
def start_assignment(self, assignment_id: UUID, node_id: UUID) -> JobAssignment:
|
|
"""The node says it has picked this up.
|
|
|
|
Guarded on the assignment already belonging to this node: without that check one
|
|
machine could start another's work by guessing an id, and the job's trail would
|
|
name a node that never touched it.
|
|
"""
|
|
|
|
now = _now()
|
|
with self._connection() as connection, self._tx(connection):
|
|
existing = connection.execute(
|
|
"SELECT node_id, state FROM public.job_assignments WHERE id = %s",
|
|
(assignment_id,),
|
|
).fetchone()
|
|
if existing is None:
|
|
raise NotFoundError("assignment_not_found")
|
|
# Ownership before state: "this belongs to someone else" and "this has already
|
|
# finished" are different problems, and an operator handed the first while
|
|
# being told the second will look in the wrong place.
|
|
if existing["node_id"] != node_id:
|
|
raise ConflictError(
|
|
"assignment_not_yours",
|
|
"This assignment belongs to another node.",
|
|
)
|
|
row = connection.execute(
|
|
"""
|
|
UPDATE public.job_assignments
|
|
SET state = 'running', started_at = COALESCE(started_at, %s)
|
|
WHERE id = %s AND node_id = %s AND state = 'assigned'
|
|
RETURNING *
|
|
""",
|
|
(now, assignment_id, node_id),
|
|
).fetchone()
|
|
if row is None:
|
|
raise ConflictError(
|
|
"assignment_not_startable",
|
|
f"This assignment is {existing['state']!r} and is not awaiting a node.",
|
|
)
|
|
connection.execute(
|
|
"UPDATE public.jobs SET status = 'running', "
|
|
"lease_owner = %s, lease_expires_at = %s "
|
|
"WHERE id = %s AND status = 'queued'",
|
|
(node_id, row["lease_expires_at"], row["job_id"]),
|
|
)
|
|
self._job_event(
|
|
connection,
|
|
row["job_id"],
|
|
"started",
|
|
{"assignment_id": str(assignment_id)},
|
|
node_id=node_id,
|
|
)
|
|
return _assignment_from_row(row)
|
|
|
|
def complete_assignment(
|
|
self,
|
|
assignment_id: UUID,
|
|
node_id: UUID,
|
|
*,
|
|
artifacts: Sequence[Mapping[str, Any]] = (),
|
|
result: Mapping[str, Any] | None = None,
|
|
) -> AssignmentOutcome:
|
|
"""The node says it finished, and hands over what it made.
|
|
|
|
What it says is not what counts. The assignment only reaches ``succeeded`` if the
|
|
controller can point at artifacts, and only then does the job reach ``complete`` -
|
|
so a node cannot mark work done by asserting it, and a lost or never-written
|
|
result is caught here rather than discovered when somebody looks for the book.
|
|
|
|
That is the rule stated as behaviour: *a job is complete when the controller has
|
|
the artifacts to show it is complete.* A node reporting success with nothing
|
|
behind it is a failure, recorded as one, with the reason naming what was missing.
|
|
|
|
Artifacts are recorded whatever the outcome. A job that failed verification may
|
|
still have produced something worth keeping - a partial render, a log - and
|
|
discarding it because the run fell short would throw away the evidence of why.
|
|
|
|
When several nodes were assigned the same job, the first verified one wins and the
|
|
rest are cancelled: their leases and heartbeats stop mattering, and leaving them
|
|
running would burn a machine's time producing a result nobody will read.
|
|
"""
|
|
|
|
now = _now()
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"SELECT * FROM public.job_assignments WHERE id = %s FOR UPDATE",
|
|
(assignment_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise NotFoundError("assignment_not_found")
|
|
if row["node_id"] != node_id:
|
|
raise ConflictError(
|
|
"assignment_not_yours",
|
|
"This assignment belongs to another node.",
|
|
)
|
|
if row["state"] not in ("assigned", "running"):
|
|
raise ConflictError(
|
|
"assignment_not_completable",
|
|
f"This assignment is {row['state']!r} and cannot be completed.",
|
|
)
|
|
|
|
job = self._job_row(connection, row["job_id"], for_update=True)
|
|
assert job is not None
|
|
|
|
recorded = self._record_artifacts(
|
|
connection,
|
|
job_id=row["job_id"],
|
|
organization_id=job["organization_id"],
|
|
project_id=job["project_id"],
|
|
node_id=node_id,
|
|
artifacts=artifacts,
|
|
)
|
|
|
|
verified, missing = self._verify_artifacts(connection, row["job_id"], job["kind"])
|
|
|
|
if not verified:
|
|
# The node said it was done and the controller disagrees. Recorded as a
|
|
# failure of the assignment, not of the job: another node may still be
|
|
# working on it, or it may be reassigned.
|
|
connection.execute(
|
|
"""
|
|
UPDATE public.job_assignments
|
|
SET state = 'failed', finished_at = %s,
|
|
error = %s, result = %s::jsonb
|
|
WHERE id = %s
|
|
""",
|
|
(
|
|
now,
|
|
f"reported complete without artifacts: missing {', '.join(missing)}",
|
|
json.dumps(dict(result or {})),
|
|
assignment_id,
|
|
),
|
|
)
|
|
self._job_event(
|
|
connection,
|
|
row["job_id"],
|
|
"unverified",
|
|
{"assignment_id": str(assignment_id), "missing": missing},
|
|
node_id=node_id,
|
|
)
|
|
self._requeue_or_fail(connection, job, node_id, now)
|
|
outcome = AssignmentOutcome(
|
|
assignment=_assignment_from_row(
|
|
self._assignment_row(connection, assignment_id)
|
|
),
|
|
verified=False,
|
|
job_status=job["status"],
|
|
missing=missing,
|
|
artifacts_recorded=recorded,
|
|
)
|
|
return outcome
|
|
|
|
finished = connection.execute(
|
|
"""
|
|
UPDATE public.job_assignments
|
|
SET state = 'succeeded', finished_at = %s,
|
|
artifacts_verified_at = %s, result = %s::jsonb
|
|
WHERE id = %s
|
|
RETURNING *
|
|
""",
|
|
(now, now, json.dumps(dict(result or {})), assignment_id),
|
|
).fetchone()
|
|
self._job_event(
|
|
connection,
|
|
row["job_id"],
|
|
"verified",
|
|
{
|
|
"assignment_id": str(assignment_id),
|
|
"artifacts": recorded,
|
|
},
|
|
node_id=node_id,
|
|
)
|
|
|
|
connection.execute(
|
|
"UPDATE public.jobs SET status = 'complete', finished_at = %s, "
|
|
"lease_owner = NULL, lease_expires_at = NULL, result = %s::jsonb "
|
|
"WHERE id = %s",
|
|
(now, json.dumps({**(result or {}), "artifacts": recorded}), row["job_id"]),
|
|
)
|
|
|
|
superseded = connection.execute(
|
|
"""
|
|
UPDATE public.job_assignments
|
|
SET state = 'cancelled', finished_at = %s,
|
|
error = 'another node completed this job first'
|
|
WHERE job_id = %s AND id <> %s
|
|
AND state IN ('assigned', 'running', 'succeeded')
|
|
RETURNING node_id
|
|
""",
|
|
(now, row["job_id"], assignment_id),
|
|
).fetchall()
|
|
|
|
self._settle_parent(
|
|
connection_organization_id=job["organization_id"],
|
|
job=_job_from_row({**job, "status": COMPLETE, "finished_at": now}),
|
|
)
|
|
|
|
return AssignmentOutcome(
|
|
assignment=_assignment_from_row(finished),
|
|
verified=True,
|
|
job_status=COMPLETE,
|
|
missing=(),
|
|
artifacts_recorded=recorded,
|
|
superseded_nodes=[str(one["node_id"]) for one in superseded],
|
|
)
|
|
|
|
def fail_assignment(
|
|
self,
|
|
assignment_id: UUID,
|
|
node_id: UUID,
|
|
*,
|
|
error: str,
|
|
retry: bool = True,
|
|
) -> AssignmentOutcome:
|
|
"""The node gives up on this assignment.
|
|
|
|
The job's fate is decided by whether anything else is still open on it. With one
|
|
node that means queued again, or `failed` once the attempts are spent; with a
|
|
fan-out the other machines are still entitled to finish it, and the job stays put.
|
|
"""
|
|
|
|
now = _now()
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"SELECT * FROM public.job_assignments WHERE id = %s FOR UPDATE",
|
|
(assignment_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise NotFoundError("assignment_not_found")
|
|
if row["node_id"] != node_id:
|
|
raise ConflictError("assignment_not_yours", "Belongs to another node.")
|
|
if row["state"] in ("succeeded", "cancelled"):
|
|
raise ConflictError(
|
|
"assignment_not_failable",
|
|
f"This assignment is {row['state']!r}.",
|
|
)
|
|
connection.execute(
|
|
"UPDATE public.job_assignments SET state = 'failed', finished_at = %s, "
|
|
"error = %s WHERE id = %s",
|
|
(now, error[:2000], assignment_id),
|
|
)
|
|
job = self._job_row(connection, row["job_id"], for_update=True)
|
|
assert job is not None
|
|
self._job_event(
|
|
connection,
|
|
row["job_id"],
|
|
"failed",
|
|
{"assignment_id": str(assignment_id), "error": error[:2000]},
|
|
node_id=node_id,
|
|
)
|
|
if not retry:
|
|
connection.execute(
|
|
"UPDATE public.jobs SET status = 'failed', finished_at = %s, "
|
|
"lease_owner = NULL, lease_expires_at = NULL, error = %s WHERE id = %s",
|
|
(now, error[:2000], row["job_id"]),
|
|
)
|
|
else:
|
|
self._requeue_or_fail(connection, job, node_id, now)
|
|
job = self._job_row(connection, row["job_id"]) or job
|
|
return AssignmentOutcome(
|
|
assignment=_assignment_from_row(self._assignment_row(connection, assignment_id)),
|
|
verified=False,
|
|
job_status=job["status"],
|
|
)
|
|
|
|
def cancel_assignment(
|
|
self,
|
|
assignment_id: UUID,
|
|
*,
|
|
reason: str,
|
|
actor: str | None = None,
|
|
) -> AssignmentOutcome:
|
|
"""An administrator stops one assignment, or every open one on a job.
|
|
|
|
Cancelling an assignment is not cancelling the job. The two differ whenever a job
|
|
was fanned out, and conflating them is how a job ends up dead because one machine
|
|
was withdrawn - when three others were still entitled to finish it.
|
|
"""
|
|
|
|
now = _now()
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"SELECT * FROM public.job_assignments WHERE id = %s FOR UPDATE",
|
|
(assignment_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise NotFoundError("assignment_not_found")
|
|
if row["state"] in ("succeeded", "cancelled"):
|
|
raise ConflictError(
|
|
"assignment_not_cancellable",
|
|
f"This assignment is {row['state']!r} and has already finished.",
|
|
)
|
|
connection.execute(
|
|
"UPDATE public.job_assignments SET state = 'cancelled', finished_at = %s, "
|
|
"error = %s WHERE id = %s",
|
|
(now, reason[:2000], assignment_id),
|
|
)
|
|
job = self._job_row(connection, row["job_id"], for_update=True)
|
|
assert job is not None
|
|
self._job_event(
|
|
connection,
|
|
row["job_id"],
|
|
"cancelled",
|
|
{"assignment_id": str(assignment_id), "reason": reason[:2000], "actor": actor},
|
|
node_id=row["node_id"],
|
|
)
|
|
self._requeue_or_fail(connection, job, row["node_id"], now)
|
|
job = self._job_row(connection, row["job_id"]) or job
|
|
return AssignmentOutcome(
|
|
assignment=_assignment_from_row(self._assignment_row(connection, assignment_id)),
|
|
verified=False,
|
|
job_status=job["status"],
|
|
)
|
|
|
|
# -- artifact verification ---------------------------------------------
|
|
|
|
def _record_artifacts(
|
|
self,
|
|
connection: Any,
|
|
*,
|
|
job_id: UUID,
|
|
organization_id: UUID,
|
|
project_id: UUID | None,
|
|
node_id: UUID,
|
|
artifacts: Sequence[Mapping[str, Any]],
|
|
) -> list[str]:
|
|
"""Write the artifacts a node says it produced, and link them to the job.
|
|
|
|
Recorded before verification and independently of it, because a run that fell
|
|
short still produced evidence and throwing it away would destroy the record of
|
|
what went wrong.
|
|
|
|
`content_hash` is required and `storage_key` with it. There is no object storage
|
|
yet (DECISION-6 is unresolved), so what is stored is the *description* of a file
|
|
and where it is meant to be - which is enough to verify that a node did the work
|
|
and to fail loudly when the bytes are not there, and not enough to pretend the
|
|
bytes exist. `storage_key` is therefore not proof of content and is never treated
|
|
as such by :meth:`_verify_artifacts`.
|
|
"""
|
|
|
|
recorded: list[str] = []
|
|
for spec in artifacts:
|
|
content_hash = str(spec.get("content_hash") or "").strip()
|
|
storage_key = str(spec.get("storage_key") or "").strip()
|
|
if not content_hash or not storage_key:
|
|
# Not stored. An artifact with neither a hash nor a key is a claim, and
|
|
# verification would then be checking the node against itself.
|
|
continue
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.artifacts (
|
|
project_id, organization_id, lang, kind, status, content_hash,
|
|
byte_size, media_type, storage_key, created_by_job_id, derivation
|
|
)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb)
|
|
RETURNING id
|
|
""",
|
|
(
|
|
project_id,
|
|
organization_id,
|
|
spec.get("lang"),
|
|
str(spec.get("kind") or "unknown"),
|
|
str(spec.get("status") or "intermediate"),
|
|
content_hash,
|
|
int(spec.get("byte_size") or 0),
|
|
str(spec.get("media_type") or "application/octet-stream"),
|
|
storage_key,
|
|
job_id,
|
|
json.dumps(dict(spec.get("derivation") or {})),
|
|
),
|
|
).fetchone()
|
|
connection.execute(
|
|
"INSERT INTO public.job_artifacts (job_id, artifact_id, role) VALUES (%s, %s, %s)",
|
|
(job_id, row["id"], str(spec.get("role") or "output")),
|
|
)
|
|
recorded.append(str(row["id"]))
|
|
return recorded
|
|
|
|
def _verify_artifacts(
|
|
self, connection: Any, job_id: UUID, kind: str
|
|
) -> tuple[bool, list[str]]:
|
|
"""Does the controller have what it needs to call this job done?
|
|
|
|
The floor is one linked artifact. `storage_key` is deliberately not consulted:
|
|
there is no object storage to ask yet, so checking it would be checking a string
|
|
the node wrote. The requirement is deliberately weak for that reason - it proves a
|
|
node produced and registered a result, not that the bytes are retrievable - and
|
|
it is stated as the floor here so that tightening it later is a change to one
|
|
method rather than a search.
|
|
|
|
`kind` is accepted and unused: what a job must produce varies by kind, and
|
|
encoding that per kind before any of them exist would be guessing. The parameter
|
|
is present so the signature does not have to change when it is.
|
|
"""
|
|
|
|
del kind
|
|
count = connection.execute(
|
|
"SELECT count(*) AS total FROM public.job_artifacts WHERE job_id = %s",
|
|
(job_id,),
|
|
).fetchone()["total"]
|
|
if count:
|
|
return True, []
|
|
return False, ["at least one artifact linked to the job"]
|
|
|
|
def _requeue_or_fail(
|
|
self, connection: Any, job: Mapping[str, Any], node_id: UUID, now: datetime
|
|
) -> None:
|
|
"""Settle a job whose holder just gave up on it.
|
|
|
|
Requeued only when nothing else is still open on it. A job with another machine
|
|
working on it is not waiting for anybody, and moving its status around underneath
|
|
that machine is how a running job gets handed out a second time.
|
|
|
|
Attempts are counted from `job_assignments`, not from `jobs.attempt`. That column
|
|
was the claim counter - `claim_job` incremented it - and under the assignment model
|
|
nothing increments it, so a job would have been re-offered forever with
|
|
`max_attempts` never reached. A job that can fail an unbounded number of times is
|
|
a job that never ends.
|
|
|
|
Counted as failed assignments rather than total assignments, so a fan-out to three
|
|
machines costs one attempt and not three: the three are the same attempt at the
|
|
same work.
|
|
"""
|
|
|
|
still_open = connection.execute(
|
|
"SELECT count(*) AS total FROM public.job_assignments "
|
|
"WHERE job_id = %s AND state IN ('assigned', 'running')",
|
|
(job["id"],),
|
|
).fetchone()["total"]
|
|
if still_open:
|
|
return
|
|
spent = connection.execute(
|
|
"SELECT count(*) AS total FROM public.job_assignments "
|
|
"WHERE job_id = %s AND state = 'failed'",
|
|
(job["id"],),
|
|
).fetchone()["total"]
|
|
if spent >= job["max_attempts"]:
|
|
connection.execute(
|
|
"UPDATE public.jobs SET status = 'failed', finished_at = %s, "
|
|
"lease_owner = NULL, lease_expires_at = NULL WHERE id = %s",
|
|
(now, job["id"]),
|
|
)
|
|
else:
|
|
connection.execute(
|
|
"UPDATE public.jobs SET status = 'queued', lease_owner = NULL, "
|
|
"lease_expires_at = NULL, not_before = now() WHERE id = %s",
|
|
(job["id"],),
|
|
)
|
|
|
|
def _candidate_nodes(
|
|
self,
|
|
connection: Any,
|
|
*,
|
|
node_ids: Sequence[UUID] | None,
|
|
kinds: Sequence[str],
|
|
) -> list[UUID]:
|
|
"""Which machines may be given this job.
|
|
|
|
The same status rule the claim path uses, applied on both routes - by name and by
|
|
kind - and for the same reason: a drained, retired or removed machine that can
|
|
still be assigned work makes the operator's controls decorative.
|
|
|
|
A named machine that is not eligible is dropped rather than refused. The caller
|
|
asked for work to be placed and it was not placed anywhere it should not have
|
|
been; an error would imply the request was malformed, which it was not, and would
|
|
make a routine drain look like a failure of the controller.
|
|
"""
|
|
|
|
if node_ids:
|
|
# Status is checked here too, and not only on the by-kind path. Naming a
|
|
# machine outright used to skip the check entirely, which meant an
|
|
# administrator could drain a node, have it appear drained in the console,
|
|
# and still hand it work - the control doing nothing while looking like it
|
|
# worked. Naming a machine is a stronger statement about *which* machine, not
|
|
# a claim that it is available.
|
|
rows = connection.execute(
|
|
"SELECT id, status, last_heartbeat_at FROM public.compute_nodes "
|
|
"WHERE id = ANY(%s) ORDER BY name",
|
|
(list(node_ids),),
|
|
).fetchall()
|
|
return [
|
|
row["id"]
|
|
for row in rows
|
|
if effective_status(row["status"], row["last_heartbeat_at"]) == "online"
|
|
]
|
|
rows = connection.execute(
|
|
"SELECT id, status, last_heartbeat_at FROM public.compute_nodes ORDER BY name"
|
|
).fetchall()
|
|
return [
|
|
row["id"]
|
|
for row in rows
|
|
if effective_status(row["status"], row["last_heartbeat_at"]) == "online"
|
|
]
|
|
|
|
def _next_attempt(self, connection: Any, job_id: UUID, node_id: UUID) -> int:
|
|
previous = connection.execute(
|
|
"SELECT max(attempt) AS highest FROM public.job_assignments "
|
|
"WHERE job_id = %s AND node_id = %s",
|
|
(job_id, node_id),
|
|
).fetchone()["highest"]
|
|
return int(previous or 0) + 1
|
|
|
|
@staticmethod
|
|
def _assignment_row(connection: Any, assignment_id: UUID) -> Mapping[str, Any]:
|
|
row = connection.execute(
|
|
"SELECT * FROM public.job_assignments WHERE id = %s", (assignment_id,)
|
|
).fetchone()
|
|
assert row is not None
|
|
return row
|
|
|
|
@staticmethod
|
|
def _job_row(
|
|
connection: Any, job_id: UUID, *, for_update: bool = False
|
|
) -> Mapping[str, Any] | None:
|
|
"""The raw job row, optionally locked.
|
|
|
|
Read as a row rather than through `get_job` because the assignment operations all
|
|
need to decide something *and* write it in one transaction. Going through the
|
|
dataclass would mean reading, closing the connection, deciding, and reopening -
|
|
and the gap between those is where two controllers assign the same job twice.
|
|
"""
|
|
|
|
query = "SELECT * FROM public.jobs WHERE id = %s"
|
|
if for_update:
|
|
query += " FOR UPDATE"
|
|
return connection.execute(query, (job_id,)).fetchone()
|
|
|
|
# -- events ------------------------------------------------------------
|
|
|
|
def _node_event(
|
|
self,
|
|
connection: Any,
|
|
node_id: UUID,
|
|
action: str,
|
|
detail: Mapping[str, Any] | None = None,
|
|
*,
|
|
actor_user_id: UUID | None = None,
|
|
) -> None:
|
|
"""Record something that happened to a node, in the audit trail.
|
|
|
|
`job_events` is the wrong home for these: its `job_id` is not nullable, so a
|
|
thing that happened to a *machine* - revoked credentials, a permanent removal -
|
|
has no job to hang from. `audit_events` is nullable in every part and has a
|
|
`target_type`/`target_id` pair, which is exactly this shape.
|
|
|
|
`actor_user_id` is nullable too, so an action taken by the controller on its own
|
|
initiative is recorded as such rather than being attributed to whoever happened
|
|
to be signed in.
|
|
"""
|
|
|
|
node = self._node_row(connection, node_id)
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO public.audit_events (
|
|
actor_user_id, organization_id, action, target_type, target_id, metadata
|
|
)
|
|
VALUES (%s, %s, %s, 'compute_node', %s, %s::jsonb)
|
|
""",
|
|
(
|
|
actor_user_id,
|
|
node["organization_id"] if node else None,
|
|
action,
|
|
str(node_id),
|
|
json.dumps(dict(detail or {})),
|
|
),
|
|
)
|
|
|
|
@staticmethod
|
|
def _job_event(
|
|
connection: Any,
|
|
job_id: UUID,
|
|
event: str,
|
|
detail: Mapping[str, Any] | None = None,
|
|
*,
|
|
node_id: UUID | None = None,
|
|
) -> None:
|
|
"""Append to a job's trail. Append-only, like every other write to that table."""
|
|
|
|
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(dict(detail or {}))),
|
|
)
|
|
|
|
@staticmethod
|
|
def _node_row(connection: Any, node_id: UUID) -> dict[str, Any] | None:
|
|
return connection.execute(
|
|
"SELECT * FROM public.compute_nodes WHERE id = %s", (node_id,)
|
|
).fetchone()
|
|
|
|
# -- 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,),
|
|
)
|
|
|
|
# -- pairing -----------------------------------------------------------
|
|
|
|
def create_pairing_code(
|
|
self,
|
|
organization_id: UUID,
|
|
node_name: str,
|
|
*,
|
|
locality: str | None = None,
|
|
trust_domain: str = "local",
|
|
ttl_seconds: int = 600,
|
|
) -> dict[str, Any]:
|
|
"""Mint a single-use code for a *named* node. Shown once, never stored.
|
|
|
|
Six digits is 1,000,000 possibilities, which is only defensible because of
|
|
what brackets it: the code is single-use, expires in ten minutes, is bound to
|
|
a node name, is Argon2id-hashed at rest, and ``PAIRING_ATTEMPT_BUDGET`` caps
|
|
guesses per source address. Any one of those alone would be too thin.
|
|
|
|
The TTL is bounded on both sides. Long enough to walk to a machine; short
|
|
enough that a code left in a terminal buffer is not a credential for the day.
|
|
"""
|
|
|
|
if ttl_seconds < 60 or ttl_seconds > 1800:
|
|
raise InputError(
|
|
"pairing_ttl_invalid",
|
|
"A pairing code must live between 60 and 1800 seconds.",
|
|
)
|
|
code = f"{secrets.randbelow(1_000_000):06d}"
|
|
expires_at = _now() + timedelta(seconds=ttl_seconds)
|
|
hint = f"ends {code[-2:]}"
|
|
with self._connection() as connection, self._tx(connection):
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.node_pairing_codes
|
|
(node_name, organization_id, locality, trust_domain,
|
|
code_hash, code_hint, expires_at)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
|
RETURNING id, expires_at
|
|
""",
|
|
(node_name, organization_id, locality, trust_domain,
|
|
self._passwords.hash_password(code), hint, expires_at),
|
|
).fetchone()
|
|
return {
|
|
"code": code,
|
|
"code_hint": hint,
|
|
"node_name": node_name,
|
|
"expires_at": row["expires_at"],
|
|
"id": row["id"],
|
|
}
|
|
|
|
def _pairing_attempt_allowed(self, connection, source_address: str | None) -> bool:
|
|
"""Charge one attempt against this address, and say whether it may proceed.
|
|
|
|
Taken *before* any code lookup, so a caller who never guesses a real code still
|
|
pays for every guess. The window resets rather than decaying, so an attacker
|
|
gets a fixed budget per interval rather than an unbounded trickle.
|
|
"""
|
|
|
|
if not source_address:
|
|
# No source address means no way to rate limit. Refusing outright would make
|
|
# the path unusable behind a proxy that does not forward one; allowing it
|
|
# silently would remove the only control that makes six digits defensible.
|
|
# Callers are expected to pass one.
|
|
return True
|
|
row = connection.execute(
|
|
"""
|
|
INSERT INTO public.node_pairing_attempts (source_address, window_started_at, attempts)
|
|
VALUES (%s, now(), 1)
|
|
ON CONFLICT (source_address) DO UPDATE
|
|
SET attempts = CASE
|
|
WHEN node_pairing_attempts.window_started_at < now() - interval '15 minutes'
|
|
THEN 1 ELSE node_pairing_attempts.attempts + 1 END,
|
|
window_started_at = CASE
|
|
WHEN node_pairing_attempts.window_started_at < now() - interval '15 minutes'
|
|
THEN now() ELSE node_pairing_attempts.window_started_at END
|
|
RETURNING attempts
|
|
""",
|
|
(source_address,),
|
|
).fetchone()
|
|
return row["attempts"] <= PAIRING_ATTEMPT_BUDGET
|
|
|
|
def redeem_pairing_code(
|
|
self,
|
|
code: str,
|
|
node_name: str,
|
|
*,
|
|
agent_version: str | None = None,
|
|
locality: str | None = None,
|
|
source_address: str | None = None,
|
|
mac: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Exchange a code for a node credential. The code is spent either way.
|
|
|
|
Verification is Argon2id over *every live code* rather than a lookup by prefix
|
|
or hint. Argon2 cannot be indexed, so this is a scan - the table holds only
|
|
codes no older than half an hour, so it is a handful of rows. The alternative,
|
|
selecting on a prefix, would make a truncated or misread code valid, which is a
|
|
far worse failure than a slow query.
|
|
|
|
Every rejection returns the same shape and, apart from the name mismatch, the
|
|
same message. A caller that can distinguish "no such code" from "wrong code"
|
|
from "expired" learns something per attempt, and the rate limiter turns that into
|
|
a workable search. The name mismatch is deliberately *not* hidden: it is not
|
|
secret, and an operator staring at a failure needs to be told the code is for a
|
|
different machine rather than left guessing.
|
|
"""
|
|
|
|
invalid = {
|
|
"ok": False,
|
|
"code": "pairing_code_invalid",
|
|
"message": (
|
|
"That pairing code is not valid. It may have expired, been already used, "
|
|
"or been mistyped. Ask for a new one."
|
|
),
|
|
}
|
|
if not code or not code.strip():
|
|
return {
|
|
"ok": False,
|
|
"code": "pairing_code_required",
|
|
"message": "A pairing code is required.",
|
|
}
|
|
with self._connection() as connection, self._tx(connection):
|
|
if not self._pairing_attempt_allowed(connection, source_address):
|
|
return {
|
|
"ok": False,
|
|
"code": "pairing_rate_limited",
|
|
"message": (
|
|
"Too many pairing attempts from this address. Wait 15 minutes, "
|
|
"or ask an administrator to pair this node from the console."
|
|
),
|
|
}
|
|
|
|
candidates = connection.execute(
|
|
"""
|
|
SELECT id, node_name, organization_id, locality, trust_domain, code_hash
|
|
FROM public.node_pairing_codes
|
|
WHERE used_at IS NULL AND expires_at > now()
|
|
ORDER BY created_at DESC
|
|
""",
|
|
).fetchall()
|
|
|
|
matched = None
|
|
for candidate in candidates:
|
|
if self._passwords.verify_password(candidate["code_hash"], code):
|
|
matched = candidate
|
|
break
|
|
|
|
if matched is None:
|
|
return invalid
|
|
|
|
if matched["node_name"] != node_name:
|
|
return {
|
|
"ok": False,
|
|
"code": "pairing_name_mismatch",
|
|
"message": (
|
|
f"This code is for a different node. This node reports as "
|
|
f"'{node_name}'; the code was issued for '{matched['node_name']}'."
|
|
),
|
|
}
|
|
|
|
# Mint the node id up front, insert the node, and only then spend the code.
|
|
#
|
|
# The order matters and the tests caught it. Spending first looks natural -
|
|
# "claim the code, then use it" - but `used_by_node_id` references
|
|
# `compute_nodes`, so claiming before inserting violates the foreign key and
|
|
# every redemption fails. The node has to exist before the code can point at
|
|
# it.
|
|
#
|
|
# Claiming is still the single-use guard: `used_at IS NULL` in the WHERE means
|
|
# two nodes racing on one code produce one winner, because the second UPDATE
|
|
# matches no rows. Because both statements are in one transaction, a loser
|
|
# rolls back its node insert too - so a race cannot leave an orphan node
|
|
# behind with no code accounting for it.
|
|
# Claim the code FIRST, and take a row lock while doing it.
|
|
#
|
|
# `SELECT ... FOR UPDATE` is the part that matters. Two nodes can verify the
|
|
# same code simultaneously - verification is a read - and both reach this
|
|
# point. Without the lock, both then try to create a node with the same name
|
|
# in the same organization, and the loser's INSERT fails on a unique
|
|
# constraint: an unhandled 500 for a condition that is entirely ordinary,
|
|
# namely two machines racing on one code.
|
|
#
|
|
# The lock serialises them. The second blocks until the first commits, then
|
|
# re-evaluates `used_at IS NULL`, matches nothing, and returns the ordinary
|
|
# invalid-code answer. No exception, no orphan node, and both callers get the
|
|
# same shape of response.
|
|
#
|
|
# The lock is taken before the node is inserted, so it is also the lock that
|
|
# orders the insert against any other redemption of the same code.
|
|
locked = connection.execute(
|
|
"""
|
|
SELECT id FROM public.node_pairing_codes
|
|
WHERE id = %s AND used_at IS NULL AND expires_at > now()
|
|
FOR UPDATE
|
|
""",
|
|
(matched["id"],),
|
|
).fetchone()
|
|
if locked is None:
|
|
# Someone else got there first.
|
|
return invalid
|
|
|
|
# A node name is unique per organization, so pairing a name that is already
|
|
# taken is an ordinary operator mistake - "add the node" run twice, or a
|
|
# retried bootstrap after a partial failure. It arrived as an unhandled
|
|
# UniqueViolation and surfaced to the caller as a bare 500, which says
|
|
# nothing about the cause. Checked first so the answer can name it.
|
|
# The final name, and the one place the naming rule is applied.
|
|
#
|
|
# `node_name` here is the *base* name: what the operator typed, and what the
|
|
# code was minted against. The suffix that makes it unique comes from the
|
|
# machine's MAC, which the controller cannot know when the code is minted -
|
|
# which is exactly why the code binds to the base name and the final name is
|
|
# settled here, where "is this name taken" can be asked of the database.
|
|
#
|
|
# With a MAC, a taken base name is disambiguated rather than refused: two
|
|
# machines called `VM238` are a normal thing and become `VM238-AC4B` and
|
|
# `VM238-3F19`. Without one there is no evidence they are different machines,
|
|
# so the refusal stands. That asymmetry is deliberate - the alternative is a
|
|
# re-run of bootstrap on an already-onboarded machine silently succeeding as
|
|
# `VM238-B`, producing a second identity for one machine.
|
|
def _taken(candidate: str) -> bool:
|
|
return (
|
|
connection.execute(
|
|
"SELECT 1 FROM public.compute_nodes "
|
|
"WHERE organization_id = %s AND name = %s",
|
|
(matched["organization_id"], candidate),
|
|
).fetchone()
|
|
is not None
|
|
)
|
|
|
|
if mac:
|
|
final_name = resolve_node_name(matched["node_name"], mac, _taken)
|
|
else:
|
|
final_name = matched["node_name"]
|
|
|
|
existing = connection.execute(
|
|
"""
|
|
SELECT id FROM public.compute_nodes
|
|
WHERE organization_id = %s AND name = %s
|
|
""",
|
|
(matched["organization_id"], final_name),
|
|
).fetchone()
|
|
if existing is not None:
|
|
return {
|
|
"ok": False,
|
|
"code": "node_name_taken",
|
|
"message": (
|
|
f"A node named '{final_name}' already exists in this "
|
|
"organization. Pairing does not replace an existing node: if you "
|
|
"are re-running a bootstrap for a node that is already onboarded, "
|
|
"its existing credential still works and no new one is needed. "
|
|
"Otherwise choose a different name, or remove the existing node "
|
|
"first."
|
|
),
|
|
}
|
|
|
|
node_id = uuid4()
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO public.compute_nodes
|
|
(id, organization_id, name, is_controller, trust_domain,
|
|
locality, transport, status, agent_version)
|
|
VALUES (%s, %s, %s, false, %s, %s, 'https', 'provisioning', %s)
|
|
""",
|
|
(node_id, matched["organization_id"], final_name,
|
|
matched["trust_domain"], locality or matched["locality"], agent_version),
|
|
)
|
|
|
|
claimed = connection.execute(
|
|
"""
|
|
UPDATE public.node_pairing_codes
|
|
SET used_at = now(), used_by_node_id = %s
|
|
WHERE id = %s
|
|
RETURNING used_by_node_id
|
|
""",
|
|
(node_id, locked["id"]),
|
|
).fetchone()
|
|
if claimed is None:
|
|
# Unreachable while the row lock holds, and kept because "unreachable"
|
|
# is not a guarantee a reader should have to take on trust.
|
|
return invalid
|
|
node_key = generate_token(32)
|
|
prefix = node_key[:8]
|
|
# A paired node gets the whole node scope set, not just enough to heartbeat.
|
|
#
|
|
# It was read-only, and that made every node write path dead. `node_agent.py`
|
|
# already calls claim, renew, complete, fail and decompose, and no credential
|
|
# pairing could ever issue was permitted any of them: the agent looked ready
|
|
# and the endpoints looked finished, and the whole surface was unusable. The
|
|
# same shape of defect as a form with no method attribute - everything
|
|
# present, nothing working.
|
|
#
|
|
# Granting it here is right because the pairing code *is* the authorization:
|
|
# minted by an authenticated administrator, single-use, expiring, bound to one
|
|
# node name, and rate-limited per source address. An administrator who mints a
|
|
# code has already decided to put that machine on the fleet.
|
|
#
|
|
# It does not weaken the boundary that matters. `NODE_SCOPES` contains nothing
|
|
# that reads project content, so a node credential cannot become a session
|
|
# token under another name. And the grant is per credential, with a node able
|
|
# to hold several, so a deliberately monitoring-only machine is issued a narrow
|
|
# credential and this one revoked - still possible, rather than made
|
|
# impossible.
|
|
paired_scopes = sorted(NODE_SCOPES)
|
|
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(node_key),
|
|
agent_version or "unknown",
|
|
paired_scopes),
|
|
)
|
|
# The plaintext key exists in this dict and nowhere else. It is not logged, not
|
|
# stored, and cannot be recovered afterwards.
|
|
return {
|
|
"ok": True,
|
|
"node_id": node_id,
|
|
"node_name": final_name,
|
|
"base_name": matched["node_name"],
|
|
"key_prefix": prefix,
|
|
"key": node_key,
|
|
"scopes": paired_scopes,
|
|
}
|
|
|
|
|
|
# -- 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 | None = None, 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 status = 'online'
|
|
ORDER BY name
|
|
""",
|
|
).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:
|
|
# Staleness is decided by `effective_status`, called from
|
|
# `_node_from_row`. This used to be a hand-rolled age check in the
|
|
# loop, which was the only place in the codebase that treated a quiet
|
|
# node as absent while the console and the heartbeat response still
|
|
# called it online. Two copies of one rule is how they drift; there is
|
|
# now one, and the loop reads it like every other caller.
|
|
#
|
|
# A node that has gone quiet must not be offered work: the job would be
|
|
# dispatched into the void and recovered by the lease sweep minutes
|
|
# later, instead of going to a machine that could have taken it at once.
|
|
node = _node_from_row(row, devices=())
|
|
if node.status != "online":
|
|
continue
|
|
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 _require_project_in_tenant(
|
|
self, connection, organization_id: UUID, project_id: UUID
|
|
) -> None:
|
|
"""Refuse a job naming a project this organization does not have.
|
|
|
|
`jobs` and `projects` now live in different schemas, so `jobs.project_id` cannot
|
|
have a foreign key - a cross-schema reference needs `postgres_fdw`. That FK was
|
|
the only thing stopping a job from naming a project belonging to a different
|
|
tenant, and migration 0005 dropped it. The check that was free in the database
|
|
has to be made here, in the layer that knows the tenant.
|
|
|
|
A real gap, recorded rather than papered over: the same applies to
|
|
`artifacts.project_id`, whose FK also went.
|
|
|
|
There is no `organization_id` predicate because after migration 0005 `projects`
|
|
has no such column: the schema it sits in *is* the tenant. So this resolves in
|
|
the caller's schema, and another tenant's project is not filtered out - it is
|
|
simply not there. That is the property the split was chosen for.
|
|
"""
|
|
|
|
# Qualified explicitly rather than relying on `search_path`. The account
|
|
# repository sets it per transaction; the compute repository does not, and
|
|
# borrowing that behaviour here would mean one more place where ambient state
|
|
# decides whether a security check runs at all. Deriving the schema from the
|
|
# organization id makes the check self-contained: it reads the caller's tenant or
|
|
# it reads nothing.
|
|
schema = tenant_schema_name(organization_id)
|
|
found = connection.execute(
|
|
sql.SQL("SELECT id FROM {}.{} WHERE id = {}").format(
|
|
sql.Identifier(schema), sql.Identifier("projects"), sql.Placeholder()
|
|
),
|
|
(project_id,),
|
|
).fetchone()
|
|
if found is None:
|
|
raise NotFoundError(
|
|
"project_not_in_organization",
|
|
f"Project {project_id} does not exist in this organization.",
|
|
)
|
|
|
|
|
|
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 project_id is not None:
|
|
self._require_project_in_tenant(connection, organization_id, project_id)
|
|
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,
|
|
*,
|
|
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.
|
|
|
|
The node's own status is checked in the same transaction, and that placement is
|
|
deliberate. ``set_node_status`` promises that draining "stops new claims and lets
|
|
running work finish", and it did not: this method never looked at the node, so a
|
|
drained machine kept being handed work by an agent that was still running, and
|
|
the operator control did nothing except change a label. Checking outside the
|
|
transaction would leave a window in which an administrator drains a node between
|
|
the check and the claim, and the one job that slips through is exactly the one
|
|
that makes the behaviour look intermittent.
|
|
|
|
``effective_status`` rather than the stored column, so a node whose agent has
|
|
died stops being offered work instead of taking a lease it will never renew.
|
|
Note what is *not* consulted: ``is_schedulable``. That also requires a device
|
|
that is not busy, and claiming a job is what makes a device busy - gating on it
|
|
here would mean no node could ever claim anything.
|
|
"""
|
|
|
|
now = _now()
|
|
expires = now + timedelta(seconds=lease_seconds)
|
|
with self._connection() as connection, self._tx(connection):
|
|
# Locked, so an administrator draining this node cannot land between the
|
|
# check below and the claim. FOR UPDATE, not SKIP LOCKED: this row is not a
|
|
# candidate to be skipped, it is the subject of the transaction.
|
|
status_row = connection.execute(
|
|
"SELECT status, last_heartbeat_at FROM public.compute_nodes "
|
|
"WHERE id = %s FOR UPDATE",
|
|
(node_id,),
|
|
).fetchone()
|
|
if status_row is None:
|
|
return None
|
|
if effective_status(
|
|
status_row["status"], status_row["last_heartbeat_at"], now=now
|
|
) != "online":
|
|
return None
|
|
|
|
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 (%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, 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, prefer_node=node_id)
|
|
if placement is None or placement.node.id != node_id:
|
|
fallback = Placement(
|
|
node=ComputeNode(
|
|
id=node_id,
|
|
organization_id=None,
|
|
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 effective_status(stored: str, last_heartbeat_at: datetime | None, *, now: datetime | None = None) -> str:
|
|
"""What a node's status is *right now*, given the clock.
|
|
|
|
The stored ``compute_nodes.status`` is a cache of the last thing anybody observed,
|
|
not the truth. It is promoted to ``online`` by a heartbeat and demoted to
|
|
``offline`` by an administrator - and nothing demotes it on the passage of time.
|
|
So a node whose agent was killed stays ``online`` in the database indefinitely,
|
|
looking schedulable to the console while nothing is listening. That is precisely
|
|
the failure invariant I24 exists to prevent: a node that has not spoken must not
|
|
appear available.
|
|
|
|
Two ways to close that, and the choice is deliberate:
|
|
|
|
- **Sweep it.** A background task demotes nodes past the stale window. This is how
|
|
expired leases are handled, so it is consistent, and it keeps the stored value
|
|
honest. But it adds a periodic task that must actually run on a deployment that
|
|
has already been taken down three boots by a stray process, and while it is not
|
|
running the stored value is wrong again.
|
|
- **Compute on read** (this). The stored value stays as-is; every read derives the
|
|
answer. A dead node is invisible the moment it dies, with no window and no
|
|
background dependency. The cost is that the database column can disagree with
|
|
what the API reports, which is why ``stored_status`` is kept on the model.
|
|
|
|
Only ``online`` is ever derived away. ``provisioning``, ``draining`` and ``retired``
|
|
are administrator decisions about intent, and intent does not expire because a
|
|
machine went quiet - demoting a draining node to offline would erase the fact that
|
|
an administrator asked for it to stop taking work.
|
|
"""
|
|
|
|
if stored != "online":
|
|
return stored
|
|
if last_heartbeat_at is None:
|
|
# Online but never heard from is a contradiction: only the heartbeat path sets
|
|
# 'online'. Treat it as absent rather than believing it.
|
|
return "offline"
|
|
moment = now or _now()
|
|
age = (moment - now_aware(last_heartbeat_at)).total_seconds()
|
|
if age > HEARTBEAT_STALE_SECONDS:
|
|
return "offline"
|
|
return "online"
|
|
|
|
|
|
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",
|
|
# Derived, not the column. See `effective_status` for why, and for what this
|
|
# deliberately does not change.
|
|
status=effective_status(row["status"], row.get("last_heartbeat_at")),
|
|
stored_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",
|
|
"AssignmentOutcome",
|
|
"ClaimedJob",
|
|
"ComputeNode",
|
|
"ComputeRepository",
|
|
"DEFAULT_LEASE_SECONDS",
|
|
"Device",
|
|
"FAILED",
|
|
"HEARTBEAT_STALE_SECONDS",
|
|
"PAIRING_ATTEMPT_BUDGET",
|
|
"effective_status",
|
|
"Job",
|
|
"ModelRequirement",
|
|
"NODE_SCOPES",
|
|
"PARTIAL",
|
|
"Placement",
|
|
"QUEUED",
|
|
"RUNNING",
|
|
"JobAssignment",
|
|
"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)
|