umbragit/UmbraTrans
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
715b7bc69f53023fc34e70e72c3d227a445862f6
UmbraTrans/umbratrans_web/compute.py
umbra audio 715b7bc69f Restore the guarantee migration 0005 had to drop 
`projects` moved into a per-tenant schema and `jobs` stayed in `public`. A foreign key
cannot cross that boundary without postgres_fdw, so 0005 dropped
`jobs_project_id_fkey`, `artifacts_project_id_fkey` and a composite `user_roles` key. The
migration recorded the consequence honestly - "the application still refuses it; the
server no longer does" - but the application did not.

So after the split nothing stopped a `public.jobs` row naming a project in a different
tenant, which is precisely the leak the schema split exists to prevent. `submit_job` now
verifies it.

Two things that were wrong in the first attempt, both caught by running the tests:

- The lookup was unqualified, relying on `search_path`. `ComputeRepository._tx` does not
  set it - only the account repository does - so the check ran against `public` and could
  never see any project, refusing every project-scoped job. The schema is now derived from
  the organization id and named explicitly: the check is self-contained, and does not
  depend on ambient state deciding whether a security check runs at all.
- It filtered on `projects.organization_id`. That column no longer exists - the schema the
  row sits in *is* the tenant, which is the whole point of the split. There is no
  predicate to write.

`test_a_nonexistent_project_is_refused` asserts that a caller cannot distinguish "no such
project" from "another tenant's project". From inside a tenant schema the two are the same
observation, and they should be.

Confirmed the tests detect the gap: with the check removed, three of five fail, including
`test_the_refused_job_is_not_created` - which exists because a refusal that still writes a
row is not a refusal.

547 passed, 29 skipped. Deployment untouched: 30 tables, 0 tenant schemas.
2026-10-03 14:16:35 -07:00

1852 lines
78 KiB
Python
Raw Blame History

"""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, Mapping, Sequence
from uuid import UUID, uuid4
import psycopg
from psycopg import sql
from .errors import ConflictError, InputError, 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 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):
row = connection.execute(
"UPDATE public.compute_nodes SET status = %s, updated_at = now() "
"WHERE id = %s RETURNING id",
(status, node_id),
).fetchone()
if row is None:
raise NotFoundError("compute_node_not_found")
node = self.get_node(node_id)
assert node is not None
return node
# -- devices -----------------------------------------------------------
def report_devices(self, node_id: UUID, devices: Sequence[Mapping[str, Any]]) -> list[Device]:
"""Replace a node's device measurements with what it just measured.
Upsert by ``(node, label)`` rather than by id, because the node's
device ids are its own and a restart with a different enumeration should
not collide with the previous run's rows.
"""
with self._connection() as connection, self._tx(connection):
for spec in devices:
connection.execute(
"""
INSERT INTO public.compute_devices
(compute_node_id, kind, label, vram_total_mb, vram_free_mb,
ram_total_mb, ram_free_mb, compute_units, throughput_factor,
supports_accelerators)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (compute_node_id, label) DO UPDATE SET
kind = EXCLUDED.kind,
vram_total_mb = EXCLUDED.vram_total_mb,
vram_free_mb = EXCLUDED.vram_free_mb,
ram_total_mb = EXCLUDED.ram_total_mb,
ram_free_mb = EXCLUDED.ram_free_mb,
compute_units = EXCLUDED.compute_units,
throughput_factor = EXCLUDED.throughput_factor,
supports_accelerators = EXCLUDED.supports_accelerators,
updated_at = now()
""",
(
node_id,
spec["kind"],
spec["label"],
int(spec.get("vram_total_mb", 0)),
int(spec.get("vram_free_mb", 0)),
int(spec.get("ram_total_mb", 0)),
int(spec.get("ram_free_mb", 0)),
spec.get("compute_units"),
float(spec.get("throughput_factor", 1.0)),
list(spec.get("supports_accelerators", ())),
),
)
node = self.get_node(node_id)
assert node is not None
return list(node.devices)
# -- credentials -------------------------------------------------------
def issue_node_credential(
self, node_id: UUID, *, agent_name: str, scopes: Sequence[str]
) -> tuple[str, str]:
"""Return ``(prefix, key)``. The key is shown once and never stored.
Scopes are validated against :data:`NODE_SCOPES` here rather than at
the call site, because a scope outside that set would be a node
credential that can read project content - which is the boundary this
table exists to hold.
"""
wanted = sorted(set(scopes))
if not wanted:
raise InputError("node_scope_required", "A node credential needs at least one scope.")
unknown = [scope for scope in wanted if scope not in NODE_SCOPES]
if unknown:
raise InputError(
"node_scope_unknown",
f"These are not node scopes: {', '.join(unknown)}. A node credential "
f"may only carry: {', '.join(sorted(NODE_SCOPES))}.",
)
key = generate_token(32)
prefix = key[:8]
with self._connection() as connection, self._tx(connection):
connection.execute(
"INSERT INTO public.node_credentials "
"(compute_node_id, key_prefix, key_hash, agent_name, scopes) "
"VALUES (%s, %s, %s, %s, %s)",
(node_id, prefix, self._passwords.hash_password(key), agent_name, wanted),
)
return prefix, key
def resolve_node_credential(self, key: str) -> tuple[UUID, UUID, list[str]] | None:
"""Turn a presented key into ``(node_id, organization_id, scopes)``.
The organization comes from the node, never from the request. That is
what makes a forged tenant identifier unnecessary rather than merely
rejected, and it is why a node needs no database credential at all.
"""
with self._connection() as connection:
row = connection.execute(
"""
SELECT c.compute_node_id, n.organization_id, c.scopes, c.key_hash, c.revoked_at
FROM public.node_credentials c
JOIN public.compute_nodes n ON n.id = c.compute_node_id
WHERE c.key_prefix = %s
""",
(key[:8],),
).fetchone()
if row is None or row["revoked_at"] is not None:
return None
if not self._passwords.verify_password(row["key_hash"], key):
return None
return row["compute_node_id"], row["organization_id"], list(row["scopes"])
def mark_credential_used(self, node_id: UUID) -> None:
with self._connection() as connection, self._tx(connection):
connection.execute(
"UPDATE public.node_credentials SET last_used_at = now() WHERE compute_node_id = %s",
(node_id,),
)
# -- 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,
) -> 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.
existing = connection.execute(
"""
SELECT id FROM public.compute_nodes
WHERE organization_id = %s AND name = %s
""",
(matched["organization_id"], matched["node_name"]),
).fetchone()
if existing is not None:
return {
"ok": False,
"code": "node_name_taken",
"message": (
f"A node named '{matched['node_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"], matched["node_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]
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",
["compute:heartbeat", "compute:read"]),
)
# 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": matched["node_name"],
"key_prefix": prefix,
"key": node_key,
"scopes": ["compute:heartbeat", "compute:read"],
}
# -- 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.
"""
now = _now()
expires = now + timedelta(seconds=lease_seconds)
with self._connection() as connection, self._tx(connection):
row = connection.execute(
"""
UPDATE public.jobs SET
status = 'running',
lease_owner = %s,
lease_expires_at = %s,
attempt = attempt + 1,
started_at = COALESCE(started_at, %s)
WHERE id IN (
SELECT j.id
FROM public.jobs j
WHERE j.status = 'queued'
AND j.not_before <= %s
AND (%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",
"ClaimedJob",
"ComputeNode",
"ComputeRepository",
"DEFAULT_LEASE_SECONDS",
"Device",
"FAILED",
"HEARTBEAT_STALE_SECONDS",
"PAIRING_ATTEMPT_BUDGET",
"effective_status",
"Job",
"ModelRequirement",
"NODE_SCOPES",
"PARTIAL",
"Placement",
"QUEUED",
"RUNNING",
"TERMINAL_STATUSES",
]
def build_compute_repository(settings, passwords=None) -> "ComputeRepository":
"""Build the compute repository from settings.
Shared by the request dependency and by the application's startup, because
the lease sweeper needs a repository before any request arrives. Building it
lazily in the sweeper instead would mean a controller restarted mid-job never
recovers that job: nothing would ever touch the compute surface, so the
repository would never be built, so the sweeper would have nothing to sweep,
so the job would sit ``running`` for ever - which is the failure the sweeper
exists to prevent, reached by the route the sweeper takes to prevent it.
"""
if not getattr(settings, "database_dsn", None):
raise RepositoryError("compute_unavailable", "No database is configured.")
return ComputeRepository(settings.database_dsn, passwords=passwords)
Reference in New Issue View Git Blame Copy Permalink