From d6cebe4d55f8b024e810e7a3234ab0c6f0602bf1 Mon Sep 17 00:00:00 2001 From: umbra audio Date: Sun, 4 Oct 2026 18:48:05 -0700 Subject: [PATCH] The controller assigns work, and a job is done only when the artifacts say so 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. --- umbratrans_web/compute.py | 736 +++++++++++++++++++++++ umbratrans_web/routes_compute.py | 273 ++++++++- umbratrans_web/tests/test_assignments.py | 567 +++++++++++++++++ umbratrans_web/tests/test_compute.py | 8 + 4 files changed, 1577 insertions(+), 7 deletions(-) create mode 100644 umbratrans_web/tests/test_assignments.py diff --git a/umbratrans_web/compute.py b/umbratrans_web/compute.py index 7d325da..2ddf9d3 100644 --- a/umbratrans_web/compute.py +++ b/umbratrans_web/compute.py @@ -99,6 +99,70 @@ 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 @@ -698,6 +762,676 @@ class ComputeRepository: 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( @@ -2201,6 +2935,7 @@ __all__ = [ "BACKOFF_BASE_SECONDS", "BLOCKED", "COMPLETE", + "AssignmentOutcome", "ClaimedJob", "ComputeNode", "ComputeRepository", @@ -2217,6 +2952,7 @@ __all__ = [ "Placement", "QUEUED", "RUNNING", + "JobAssignment", "TERMINAL_STATUSES", ] diff --git a/umbratrans_web/routes_compute.py b/umbratrans_web/routes_compute.py index 02526d6..8db5d5b 100644 --- a/umbratrans_web/routes_compute.py +++ b/umbratrans_web/routes_compute.py @@ -26,7 +26,7 @@ from __future__ import annotations from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, Literal from uuid import UUID from fastapi import APIRouter, Depends, Header, Query, Request, Response @@ -724,19 +724,87 @@ def set_node_status( return _node_response(repository.set_node_status(node_id, payload.status)) -class RemoveNodeRequest(RequestModel): - """Why a machine is being taken out, and who did it. +class ReasonRequest(RequestModel): + """Why, and who. Both recorded. - Both are required of the caller rather than defaulted. The reason ends up in the - audit trail next to the revocation, and a removal with no recorded cause is the one - thing that makes a later "why is this machine gone" unanswerable - dead hardware and - a decommissioned licence look identical six months afterwards. + A cancellation with no recorded cause is the one thing that makes "why did this job + stop" unanswerable a month later, when the difference between a withdrawn machine and + a superseded one is exactly what an operator is trying to establish. """ reason: str = Field(min_length=1, max_length=500) actor: str | None = Field(default=None, max_length=120) +class RemoveNodeRequest(ReasonRequest): + """Removal takes the same fields; the class exists so the two read differently. + + Both are required rather than defaulted for the same reason: dead hardware and a + decommissioned licence look identical six months afterwards. + """ + + +class AssignmentResponse(BaseModel): + """One machine's attempt at one job, as the controller sees it. + + ``artifacts_verified_at`` is the field that matters and it is deliberately not + derivable from ``state``: a ``succeeded`` assignment whose timestamp is null is a node + that says it finished and has not been confirmed, and the two must not look alike. + """ + + id: UUID + job_id: UUID + node_id: UUID + state: str + attempt: int + 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 + verified: bool + error: str | None = None + + +class AssignmentOutcomeResponse(BaseModel): + """The result of an operation on an assignment, including what verification wanted.""" + + assignment: AssignmentResponse + verified: bool + job_status: str + missing: list[str] = Field(default_factory=list) + artifacts_recorded: list[str] = Field(default_factory=list) + superseded_nodes: list[str] = Field(default_factory=list) + + +def _assignment_response(assignment: Any) -> AssignmentResponse: + return AssignmentResponse( + id=assignment.id, + job_id=assignment.job_id, + node_id=assignment.node_id, + state=assignment.state, + attempt=assignment.attempt, + assigned_at=assignment.assigned_at, + started_at=assignment.started_at, + finished_at=assignment.finished_at, + lease_expires_at=assignment.lease_expires_at, + artifacts_verified_at=assignment.artifacts_verified_at, + verified=assignment.is_verified, + error=assignment.error, + ) + + +def _outcome_response(outcome: Any) -> AssignmentOutcomeResponse: + return AssignmentOutcomeResponse( + assignment=_assignment_response(outcome.assignment), + verified=outcome.verified, + job_status=outcome.job_status, + missing=list(outcome.missing), + artifacts_recorded=list(outcome.artifacts_recorded), + superseded_nodes=list(outcome.superseded_nodes), + ) + + @admin_router.post("/nodes/{node_id}/revoke", response_model=NodeResponse) def revoke_node_credentials( node_id: UUID, @@ -787,6 +855,197 @@ def remove_node( return _node_response(repository.remove_node(node_id, reason=payload.reason, actor=payload.actor)) +class AssignJobRequest(RequestModel): + """Who should get this job, and how many of them. + + ``fan_out`` defaults to 1 and is never inferred. More than one machine means the same + work done several times over and the first verified result winning, which is worth + having for cheap idempotent work and wasteful for anything that costs GPU-seconds. + """ + + node_ids: list[UUID] | None = None + fan_out: int = Field(default=1, ge=1, le=16) + lease_seconds: int = Field(default=60, ge=10, le=3600) + + +class ArtifactSpec(RequestModel): + """One file a node says it produced. + + Every bound here mirrors a database constraint, deliberately. ``content_hash`` is 64 + lowercase hex characters because `artifacts_content_hash_check` says so, ``kind`` is + 64 because `artifacts_kind_check` says so, and ``status`` is an enum because + `artifacts_status_check` says so. Matching them here rather than discovering them at + the INSERT is the difference between a node getting a 400 it can act on and a 500 it + cannot. + + ``content_hash`` and ``storage_key`` are both required, and an artifact without them + is not stored at all - verification would otherwise be checking the node's report + against the same node's report, which is not a check. + + There is no object storage behind this yet (DECISION-6 is unresolved), so what is + recorded is the description of a file and where it is meant to be. That is enough to + establish that a node did the work, and not enough to pretend the bytes are + retrievable, so nothing here treats a ``storage_key`` as proof that it is. + """ + + kind: str = Field(min_length=1, max_length=64) + content_hash: str = Field(pattern=r"^[0-9a-f]{64}$") + storage_key: str = Field(min_length=1, max_length=1024) + byte_size: int = Field(default=0, ge=0) + media_type: str = Field(default="application/octet-stream", max_length=200) + role: Literal["produced", "consumed"] = "produced" + """`job_artifacts_role_check` allows exactly these two, and nothing else. + + A node reporting what it made is `produced`. Worth being explicit that this is a + vocabulary the database fixed and the API had to learn rather than invent: the first + version of this field defaulted to `output`, which no constraint mentioned, so every + single completion would have failed with a check violation at the INSERT. + """ + lang: str | None = Field(default=None, max_length=16) + status: Literal["intermediate", "final", "superseded"] = "intermediate" + derivation: dict[str, Any] = Field(default_factory=dict) + + +class CompleteAssignmentRequest(RequestModel): + artifacts: list[ArtifactSpec] = Field(default_factory=list) + result: dict[str, Any] = Field(default_factory=dict) + + +class FailAssignmentRequest(RequestModel): + error: str = Field(min_length=1, max_length=2000) + retry: bool = True + + +@admin_router.post("/jobs/{job_id}/assign", response_model=list[AssignmentResponse]) +def assign_job( + job_id: UUID, + payload: AssignJobRequest, + request: Request, + principal: Principal = Depends(require_admin), +) -> list[AssignmentResponse]: + """Hand a job to one or more nodes. + + An empty list is not an error: it means no machine was eligible, which is a fact an + operator needs to see rather than a 404. A job nobody can run is a queue that is not + draining, and it should say so. + """ + + repository = compute_repository(request) + assignments = repository.assign_job( + job_id, + node_ids=payload.node_ids, + fan_out=payload.fan_out, + lease_seconds=payload.lease_seconds, + actor=getattr(principal, "email", None), + ) + return [_assignment_response(one) for one in assignments] + + +@admin_router.post("/assignments/{assignment_id}/cancel", response_model=AssignmentOutcomeResponse) +def cancel_assignment( + assignment_id: UUID, + payload: ReasonRequest, + request: Request, + principal: Principal = Depends(require_admin), +) -> AssignmentOutcomeResponse: + """Stop one attempt. The job itself carries on if anything else is open on it.""" + + repository = compute_repository(request) + outcome = repository.cancel_assignment( + assignment_id, reason=payload.reason, actor=getattr(principal, "email", None) + ) + return _outcome_response(outcome) + + +@admin_router.get("/jobs/{job_id}/assignments", response_model=list[AssignmentResponse]) +def list_job_assignments( + job_id: UUID, + request: Request, + principal: Principal = Depends(require_admin), +) -> list[AssignmentResponse]: + """Every attempt made on this job, and which of them the controller verified. + + This is the answer to "what was assigned, and what is not complete". + """ + + repository = compute_repository(request) + return [_assignment_response(one) for one in repository.assignments_for_job(job_id)] + + +@agent_router.get("/assignments", response_model=list[AssignmentResponse]) +def node_assignments( + request: Request, + node: NodeContext = Depends(require_node), + limit: int = Query(default=10, ge=1, le=100), +) -> list[AssignmentResponse]: + """What this node has been asked to do. + + Polling rather than pushed, and that 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. + An agent that polls in a loop gets assignments as an ordinary consequence of running. + """ + + node.require("compute:claim") + repository = compute_repository(request) + return [ + _assignment_response(one) + for one in repository.assignments_for_node(node.node_id, limit=limit) + ] + + +@agent_router.post("/assignments/{assignment_id}/start", response_model=AssignmentResponse) +def start_assignment( + assignment_id: UUID, + request: Request, + node: NodeContext = Depends(require_node), +) -> AssignmentResponse: + node.require("compute:claim") + repository = compute_repository(request) + return _assignment_response(repository.start_assignment(assignment_id, node.node_id)) + + +@agent_router.post( + "/assignments/{assignment_id}/complete", response_model=AssignmentOutcomeResponse +) +def complete_assignment( + assignment_id: UUID, + payload: CompleteAssignmentRequest, + request: Request, + node: NodeContext = Depends(require_node), +) -> AssignmentOutcomeResponse: + """Report the work finished, and hand over what was produced. + + Reporting is not believing. The job only reaches ``complete`` if the controller can + point at artifacts, so a node that says it finished with nothing behind it gets its + assignment failed and the reason naming what was absent. + """ + + node.require("compute:complete") + repository = compute_repository(request) + outcome = repository.complete_assignment( + assignment_id, + node.node_id, + artifacts=[spec.model_dump() for spec in payload.artifacts], + result=payload.result, + ) + return _outcome_response(outcome) + + +@agent_router.post("/assignments/{assignment_id}/fail", response_model=AssignmentOutcomeResponse) +def fail_assignment( + assignment_id: UUID, + payload: FailAssignmentRequest, + request: Request, + node: NodeContext = Depends(require_node), +) -> AssignmentOutcomeResponse: + node.require("compute:complete") + repository = compute_repository(request) + outcome = repository.fail_assignment( + assignment_id, node.node_id, error=payload.error, retry=payload.retry + ) + return _outcome_response(outcome) + + @admin_router.get("/models", response_model=list[ModelResponse]) def list_models(request: Request, principal: Principal = Depends(require_admin)) -> list[ModelResponse]: """The global model registry. diff --git a/umbratrans_web/tests/test_assignments.py b/umbratrans_web/tests/test_assignments.py new file mode 100644 index 0000000..af61055 --- /dev/null +++ b/umbratrans_web/tests/test_assignments.py @@ -0,0 +1,567 @@ +"""The controller assigns work, and a job is done when the controller can prove it. + +Two rules, and the second is the one with teeth: + +* **The controller assigns.** A node no longer claims work; it 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. + +* **A job is complete when the controller has the artifacts to show it.** Not when a node + says so. A node reporting success with nothing behind it is a failure, recorded as one, + with the reason naming what was missing. + +The second rule is easy to state and easy to implement wrongly: the tempting version +marks the job complete inside the same transaction that accepted the node's report, and +then verifies afterwards "as a check". That version cannot fail, because the only thing it +could have compared against is the report it just trusted. +""" + +from __future__ import annotations + +import os +from uuid import uuid4 + +import pytest + +from umbratrans_web.compute import COMPLETE, FAILED, QUEUED, RUNNING, ComputeRepository +from umbratrans_web.tests.test_compute import _device, _node, _online + +INTEGRATION_ENV = "UMBRATRANS_TEST_DSN" + + +def _integration_enabled() -> bool: + return bool(os.environ.get(INTEGRATION_ENV)) and os.environ.get( + "UMBRATRANS_TEST_DATABASE_IS_DISPOSABLE", "" + ) in {"1", "true", "yes", "on"} + + +integration = pytest.mark.skipif( + not _integration_enabled(), + reason=f"needs {INTEGRATION_ENV} and an acknowledged disposable database", +) + + +@pytest.fixture +def compute(): + """A repository against the scratch database, emptied on teardown. + + Duplicated from `test_compute.py` rather than moved to `conftest.py`: those fixtures + are module-local on purpose, and hoisting them would make every file's fixtures depend + on one file's imports. The cost is that a change to the teardown list has to be made + twice, which is the smaller of the two mistakes. + """ + + from umbratrans_web.tests.test_compute import _require_scratch_database + + _require_scratch_database() + repository = ComputeRepository( + os.environ[INTEGRATION_ENV], min_size=1, max_size=6, passwords=None + ) + try: + yield repository + with repository._connection() as connection, connection.transaction(): # noqa: SLF001 + connection.execute( + "TRUNCATE public.job_events, public.jobs, public.job_assignments, " + "public.job_artifacts, public.artifacts, public.node_heartbeats, " + "public.node_models, public.node_credentials, public.compute_devices, " + "public.compute_nodes, public.model_registry CASCADE" + ) + finally: + repository.close() + + +@pytest.fixture +def organization(compute): + """One organization, seeded directly. + + Created in the repository rather than through the service because these tests are + about dispatch, and the service path would need an authenticated administrator to + exist first. + """ + + identifier = uuid4() + with compute._connection() as connection, connection.transaction(): # noqa: SLF001 + connection.execute( + "INSERT INTO public.organizations (id, name, slug) VALUES (%s, %s, %s)", + (identifier, "Assignment Test", f"assign-{identifier.hex[:8]}"), + ) + return identifier + + +@pytest.fixture +def world(compute, organization): + """Two online machines, a device each, and a project the jobs belong to. + + The project is not incidental. `artifacts.project_id` is NOT NULL while + `jobs.project_id` is nullable, so a job with no project cannot produce an artifact and + therefore cannot be completed - every test here would be measuring a constraint + violation instead of the thing under test. Which is itself worth knowing: work that + produces output has to belong to something. + """ + + from umbratrans_web.tenancy import tenant_search_path + + project = uuid4() + with compute._connection() as connection, compute._tx(connection): + connection.execute( + "SELECT set_config('search_path', %s, true)", (tenant_search_path(organization),) + ) + connection.execute( + "INSERT INTO projects (id, name, slug) VALUES (%s, %s, %s)", + (project, "Assigned Book", f"assigned-{project.hex[:8]}"), + ) + + first = _online( + compute, + _node(compute, organization, "assign-a"), + [_device(_node(compute, organization, "device-a"))], + ) + second = _online( + compute, + _node(compute, organization, "assign-b"), + [_device(_node(compute, organization, "device-b"))], + ) + return { + "compute": compute, + "organization": organization, + "project": project, + "a": first, + "b": second, + } + + +def _job(world, **overrides): + """A job that belongs to the project's organization, so its artifacts can exist.""" + + kwargs = {"kind": "speak", "requested_by": "test", "project_id": world["project"]} + kwargs.update(overrides) + return world["compute"].submit_job(world["organization"], **kwargs) + + +def _artifact(**overrides) -> dict: + """A well-formed artifact. Everything verification needs is present by default.""" + + spec = { + "kind": "audio", + "content_hash": "a" * 64, + "storage_key": "audio/chapter-001.mp3", + "byte_size": 1024, + "media_type": "audio/mpeg", + "role": "produced", + } + spec.update(overrides) + return spec + + +# -- the controller assigns --------------------------------------------------- + + +@integration +def test_a_job_goes_to_the_node_the_controller_named(world) -> None: + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + + created = compute.assign_job(job.id, node_ids=[a.id]) + + assert len(created) == 1 + assert created[0].node_id == a.id + assert created[0].state == "assigned" + + +@integration +def test_a_job_is_not_running_until_a_node_starts_it(world) -> None: + """Assignment is not progress. + + A console that says ``running`` for a job nobody has picked up yet is claiming work + that has not happened, and the difference matters when an operator is deciding whether + the queue is moving. + """ + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + + compute.assign_job(job.id, node_ids=[a.id]) + + assert compute.get_job(job.id).status == QUEUED, "an unstarted job claims to be running" + + +@integration +def test_a_node_sees_only_what_was_assigned_to_it(world) -> None: + """A node's whole view of the queue, and it is only its own work.""" + + compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"] + mine = _job(world) + theirs = _job(world) + compute.assign_job(mine.id, node_ids=[a.id]) + compute.assign_job(theirs.id, node_ids=[b.id]) + + seen_by_a = compute.assignments_for_node(a.id) + + assert [one.job_id for one in seen_by_a] == [mine.id], ( + f"node A was shown work that is not its own: {[str(one.job_id) for one in seen_by_a]}" + ) + + +@integration +def test_polling_does_not_hand_back_work_already_in_hand(world) -> None: + """An agent polls in a loop. Without this it would start the same job every time.""" + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + + assert compute.assignments_for_node(a.id) == [], ( + "a running assignment is offered again to the node already running it" + ) + + +@integration +def test_work_comes_in_priority_order(world) -> None: + """The controller expresses urgency through priority, so the node must honour it. + + Ordered by assignment time instead, a node picks up whatever arrived first and the + priority column becomes decoration. + """ + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + low = _job(world, priority=900) + high = _job(world, priority=1) + compute.assign_job(low.id, node_ids=[a.id]) + compute.assign_job(high.id, node_ids=[a.id]) + + seen = compute.assignments_for_node(a.id) + + assert [one.job_id for one in seen] == [high.id, low.id], ( + "the node was given the low-priority work first" + ) + + +@integration +def test_a_drained_node_is_not_given_work_by_assignment(world) -> None: + """The claim path checks status; the assignment path must too, or draining is theatre. + + Draining a machine and then assigning to it by name would be the worst of both: the + operator's control appears to work, and a node explicitly switched off keeps receiving + work. + """ + + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + compute.set_node_status(a.id, "draining") + + created = compute.assign_job(job.id, node_ids=[a.id]) + + assert not any(one.node_id == a.id for one in created), ( + "a drained node was assigned work" + ) + assert compute.assignments_for_job(job.id) == [] + + +@integration +def test_one_node_cannot_start_another_nodes_assignment(world) -> None: + """Guessing an id must not put a machine's name on work it never did.""" + + from umbratrans_web.errors import ConflictError + + compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + + with pytest.raises(ConflictError) as error: + compute.start_assignment(assignment.id, b.id) + + assert error.value.code == "assignment_not_yours", ( + f"the refusal does not say the assignment is another node's: {error.value}" + ) + assert compute.assignments_for_job(job.id)[0].state == "assigned" + + +# -- a job is complete when the controller has the artifacts ------------------ + + +@integration +def test_completing_with_artifacts_completes_the_job(world) -> None: + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + + outcome = compute.complete_assignment( + assignment.id, a.id, artifacts=[_artifact()], result={"chapters": 1} + ) + + assert outcome.verified is True + assert outcome.job_status == COMPLETE + assert compute.get_job(job.id).status == COMPLETE + assert outcome.assignment.artifacts_verified_at is not None + + +@integration +def test_a_node_reporting_success_with_nothing_behind_it_does_not_complete_the_job(world) -> None: + """The rule with teeth. + + The node says it finished. The controller looks for artifacts, finds none, and records + a failure - because "a job is complete when the controller has the artifacts to show + it" means the node's report is a claim to be checked, not a fact to be recorded. + """ + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + + outcome = compute.complete_assignment(assignment.id, a.id, artifacts=[], result={"chapters": 1}) + + assert outcome.verified is False + assert compute.get_job(job.id).status != COMPLETE, ( + "the job was completed on the node's word alone, with no artifact behind it" + ) + assert outcome.missing, "the failure does not say what was missing" + + +@integration +def test_an_artifact_without_a_hash_is_not_stored_and_does_not_count(world) -> None: + """Otherwise verification would be checking the node's report against itself. + + An artifact with neither a hash nor a storage key is a claim. Accepting it would mean + the verification step compares a node's report with the same node's report, which is + not a check at all - and it would pass. + """ + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + + outcome = compute.complete_assignment( + assignment.id, a.id, artifacts=[{"kind": "audio", "byte_size": 10}] + ) + + assert outcome.verified is False, "a hashless artifact was accepted as proof" + assert outcome.artifacts_recorded == [] + + +@integration +def test_verification_is_recorded_separately_from_the_node_saying_so(world) -> None: + """`succeeded` and `verified` are different facts and must not look alike. + + If they collapse, then a node's report is the definition of done, and the column that + records independent confirmation has nothing to record. + """ + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + compute.complete_assignment(assignment.id, a.id, artifacts=[_artifact()]) + + stored = compute.assignments_for_job(job.id)[0] + + assert stored.state == "succeeded" + assert stored.is_verified is True + assert stored.artifacts_verified_at is not None + + +@integration +def test_the_artifact_is_linked_to_the_job_so_the_answer_is_readable(world) -> None: + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + outcome = compute.complete_assignment(assignment.id, a.id, artifacts=[_artifact()]) + + assert len(outcome.artifacts_recorded) == 1 + with compute._connection() as connection: + linked = connection.execute( + "SELECT count(*) AS total FROM public.job_artifacts WHERE job_id = %s", (job.id,) + ).fetchone()["total"] + assert linked == 1, "the artifact is not linked to the job, so nothing records what it produced" + + +# -- more than one node ------------------------------------------------------- + + +@integration +def test_a_job_can_be_given_to_several_machines_at_once(world) -> None: + """What `lease_owner` could not express.""" + + compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + + created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2) + + assert {one.node_id for one in created} == {a.id, b.id} + assert len(compute.assignments_for_job(job.id)) == 2 + + +@integration +def test_one_machine_is_not_given_the_same_job_twice(world) -> None: + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + + first = compute.assign_job(job.id, node_ids=[a.id]) + second = compute.assign_job(job.id, node_ids=[a.id]) + + assert len(first) == 1 + assert len(compute.assignments_for_job(job.id)) == 1, "the same node holds the job twice" + + +@integration +def test_the_first_verified_machine_wins_and_the_others_are_cancelled(world) -> None: + """Otherwise every machine on the list burns its time on a result nobody reads.""" + + compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2) + by_node = {one.node_id: one for one in created} + + compute.start_assignment(by_node[a.id].id, a.id) + outcome = compute.complete_assignment(by_node[a.id].id, a.id, artifacts=[_artifact()]) + + assert outcome.verified is True + assert str(b.id) in outcome.superseded_nodes + loser = next( + one for one in compute.assignments_for_job(job.id) if one.node_id == b.id + ) + assert loser.state == "cancelled", "the losing machine is still holding an open assignment" + assert compute.get_job(job.id).status == COMPLETE + + +@integration +def test_a_machine_giving_up_does_not_disturb_one_still_working(world) -> None: + """With a fan-out the job is not waiting for anybody. + + Requeuing or failing the job while another machine is mid-flight is how a running job + gets handed out a second time. + """ + + compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2) + by_node = {one.node_id: one for one in created} + compute.start_assignment(by_node[b.id].id, b.id) + compute.start_assignment(by_node[a.id].id, a.id) + + compute.fail_assignment(by_node[a.id].id, a.id, error="gpu caught fire") + + assert compute.get_job(job.id).status == RUNNING, ( + "the job was moved out from under the machine that is still working on it" + ) + assert next( + one for one in compute.assignments_for_job(job.id) if one.node_id == b.id + ).state == "running" + + +@integration +def test_when_nothing_is_left_open_the_work_is_offered_again(world) -> None: + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + + compute.fail_assignment(assignment.id, a.id, error="out of memory") + + assert compute.get_job(job.id).status == QUEUED, ( + "the job is neither running nor waiting for anybody, so nothing will pick it up" + ) + + +@integration +def test_work_is_offered_again_until_the_attempts_run_out(world) -> None: + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world, max_attempts=1) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + + compute.fail_assignment(assignment.id, a.id, error="out of memory") + + assert compute.get_job(job.id).status == FAILED, ( + "the job keeps being re-offered after its attempts are spent, so it never ends" + ) + + +# -- cancelling --------------------------------------------------------------- + + +@integration +def test_cancelling_one_attempt_leaves_the_job_for_the_others(world) -> None: + """Cancelling an assignment is not cancelling the job. + + Conflating them is how a job dies because one machine was withdrawn while three + others were still entitled to finish it. + """ + + compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + created = compute.assign_job(job.id, node_ids=[a.id, b.id], fan_out=2) + by_node = {one.node_id: one for one in created} + + compute.cancel_assignment(by_node[a.id].id, reason="machine withdrawn") + + assert compute.get_job(job.id).status != FAILED + assert next( + one for one in compute.assignments_for_job(job.id) if one.node_id == b.id + ).state == "assigned", "the other machine lost its assignment when one was cancelled" + + +@integration +def test_cancelling_records_why(world) -> None: + """A month later, withdrawn and superseded look identical unless something says so.""" + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + + compute.cancel_assignment(assignment.id, reason="machine withdrawn for reimage", actor="ops") + + events = [event["event"] for event in compute.job_events(job.id)] + assert "cancelled" in events + details = [ + event["detail"] + for event in compute.job_events(job.id) + if event["event"] == "cancelled" + ] + assert any("reimage" in str(detail) for detail in details), ( + f"the cancellation was recorded without its reason: {details}" + ) + + +@integration +def test_a_finished_job_cannot_be_assigned(world) -> None: + """Handing completed work to a machine is duplicated effort with a real cost.""" + + from umbratrans_web.errors import ConflictError + + compute, organization, a, _b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + compute.complete_assignment(assignment.id, a.id, artifacts=[_artifact()]) + + with pytest.raises(ConflictError): + compute.assign_job(job.id, node_ids=[a.id]) + + +@integration +def test_a_failed_verification_leaves_the_work_available(world) -> None: + """Not lost, and not completed: available for another machine to be given. + + A node reporting success with no artifact is a failed *attempt*. If that finished the + job, one broken node would silently end a queue. + """ + + compute, organization, a, b = world["compute"], world["organization"], world["a"], world["b"] + job = _job(world) + assignment = compute.assign_job(job.id, node_ids=[a.id])[0] + compute.start_assignment(assignment.id, a.id) + + compute.complete_assignment(assignment.id, a.id, artifacts=[]) + + assert compute.get_job(job.id).status == QUEUED + retry = compute.assign_job(job.id, node_ids=[b.id]) + assert len(retry) == 1, "the work cannot be given to another machine after a failed attempt" + + diff --git a/umbratrans_web/tests/test_compute.py b/umbratrans_web/tests/test_compute.py index 564c800..b2c18f5 100644 --- a/umbratrans_web/tests/test_compute.py +++ b/umbratrans_web/tests/test_compute.py @@ -1171,6 +1171,14 @@ def test_the_node_surface_is_exactly_what_a_node_may_do() -> None: "/api/node/bootstrap", "/api/node/heartbeat", "/api/node/models/{slug}/state", + # The controller assigns work; the node asks what it has been given. These + # replace the claim endpoints below, which are still here only until the + # migration off them is finished - see test_claim_is_still_reachable_only_as_a + # _deprecated_path, which is what will fail when they are removed. + "/api/node/assignments", + "/api/node/assignments/{assignment_id}/start", + "/api/node/assignments/{assignment_id}/complete", + "/api/node/assignments/{assignment_id}/fail", "/api/node/jobs/claim", "/api/node/jobs/{job_id}/renew", "/api/node/jobs/{job_id}/complete",