0f5c1e3408
The compute proof of concept had never been run with the agent as umbra-poc. Running it that way found three defects, all the same shape: a check that passed because its mechanism did nothing. The demonstration never reached the agent at all. It held the pid it spawned, which under `sudo -n -u` is the sudo wrapper - sudo forks the agent before exec'ing. Signalling that killed the wrapper and left the worker running, renewing its lease and holding the job indefinitely, under a printed note saying "exactly like a power cut". An unprivileged process cannot signal another account's process at all, and kill(2) on a process group succeeds as soon as one member is signalled, so the wrapper's death read as the agent's. Every run leaked an agent. Lease recovery was a diagnosis rather than a mechanism. sweep_expired_leases documented itself as "called by the controller's own timer" and there was no timer; a killed node's job came back only when somebody noticed a stalled job and posted to an administrative route by hand. The proof of concept swept by hand and recorded the result as automatic. app.py now runs the loop for the life of the application, building the compute repository at startup so a restart mid-job still recovers it. The power cut is now the agent's own: os._exit(137) the instant it claims a job, before running it, because a process that does not own the agent cannot signal it. A clean stop goes through a stop_file the agent polls, for the same reason. Liveness is measured by finding the worker's own pid in /proc matched on its config path, sampled repeatedly and accumulated - the wrapper and the worker are not simultaneous, so a single sample finds only the wrapper. The node key travels by `sudo --preserve-env` rather than as VAR=value in argv, because argv is world-readable through /proc/<pid>/cmdline and the environment is not. The demonstration checks whether the key will survive before starting and falls back to a pipe, which node_agent.py --node-key-stdin accepts. Neither a file nor argv ever holds it. Eight tests, each verified to fail against a mutation of the code it covers: three for the sweeper (it must return an expired lease, must survive a failing round, must not invent a repository it has no database for) and five driving a real agent process against a stub controller (stops on the stop file, exits 137 holding an uncompleted job, takes the key from both routes and neither, reports an empty pipe as an empty pipe). 448 pass with a database, 421 without. The deployment database is untouched: users=2 orgs=3 projects=0 audit=43, no compute tables.
979 lines
40 KiB
Python
979 lines
40 KiB
Python
"""Compute substrate tests.
|
|
|
|
Four of these are about claims and leases, and they exist because the alternative
|
|
is a failure that is silent when it happens:
|
|
|
|
* **I25** two agents claiming at once must get different jobs. A queue where they
|
|
can get the same one produces duplicated work that is individually plausible
|
|
and collectively wrong.
|
|
* **I26** a job whose holder was killed must come back. A worker that is
|
|
``kill -9``'d releases nothing, so the lease is what releases it.
|
|
* **I24** a node that has never authenticated must not be schedulable. A node that
|
|
looks ready before its agent has spoken swallows every job dispatched to it.
|
|
* **I28** a node credential must not reach the administrator surface. That one is
|
|
here as a route-list assertion, because it was a real bug: mounting one router
|
|
at two prefixes published node creation behind node-credential authentication.
|
|
|
|
And one about the thing that is easiest to get wrong in the other direction:
|
|
|
|
* a slow device is **eligible**. A CPU with enough memory takes the job; refusing
|
|
it would turn a machine that would have finished the work in an hour into a
|
|
machine that never touches it. §7.3
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import concurrent.futures
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from uuid import UUID, uuid4
|
|
|
|
import pytest
|
|
|
|
from umbratrans_web.compute import (
|
|
BLOCKED,
|
|
COMPLETE,
|
|
FAILED,
|
|
PARTIAL,
|
|
QUEUED,
|
|
RUNNING,
|
|
ComputeRepository,
|
|
ModelRequirement,
|
|
)
|
|
|
|
INTEGRATION_ENV = "UMBRATRANS_TEST_DSN"
|
|
ACK_ENV = "UMBRATRANS_TEST_DATABASE_IS_DISPOSABLE"
|
|
|
|
integration = pytest.mark.skipif(
|
|
not os.environ.get(INTEGRATION_ENV) or os.environ.get(ACK_ENV) != "1",
|
|
reason=(
|
|
f"{INTEGRATION_ENV} is unset, or {ACK_ENV} is not 1. This suite writes to and "
|
|
"truncates tables in the target database, so it needs a scratch database and an "
|
|
"explicit acknowledgement. Never point it at a deployment."
|
|
),
|
|
)
|
|
|
|
|
|
def _require_scratch_database() -> None:
|
|
"""Refuse to run without a scratch database and an explicit acknowledgement.
|
|
|
|
Repeated inside the fixture rather than relying on the ``integration`` mark,
|
|
because a fixture cannot take a mark as an argument - and a skipped
|
|
*test* with a fixture that still ran would truncate a database before
|
|
skipping.
|
|
"""
|
|
|
|
if not os.environ.get(INTEGRATION_ENV) or os.environ.get(ACK_ENV) != "1":
|
|
pytest.skip(
|
|
f"{INTEGRATION_ENV} is unset, or {ACK_ENV} is not 1. This suite truncates "
|
|
"tables in the target database, so it needs a scratch database and an explicit "
|
|
"acknowledgement. Never point it at a deployment."
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def compute():
|
|
"""A repository against a scratch database, emptied on teardown."""
|
|
|
|
_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.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, "Compute Test", f"compute-{identifier.hex[:8]}"),
|
|
)
|
|
return identifier
|
|
|
|
|
|
def _node(compute, organization_id: UUID, name: str, **kwargs):
|
|
return compute.create_node(organization_id, name=name, **kwargs)
|
|
|
|
|
|
def _device(node, *, kind="cpu", vram_total=0, vram_free=0, ram_total=16000, ram_free=8000, factor=1.0, **extra):
|
|
return {
|
|
"kind": kind,
|
|
"label": extra.get("label", "cpu:0" if kind == "cpu" else "gpu:0"),
|
|
"vram_total_mb": vram_total,
|
|
"vram_free_mb": vram_free,
|
|
"ram_total_mb": ram_total,
|
|
"ram_free_mb": ram_free,
|
|
"compute_units": extra.get("compute_units"),
|
|
"throughput_factor": factor,
|
|
"supports_accelerators": extra.get("supports_accelerators", []),
|
|
}
|
|
|
|
|
|
def _online(compute, node, devices):
|
|
"""Bring a node to ``online`` the only way it can become online: a heartbeat."""
|
|
|
|
return compute.heartbeat(node.id, agent_version="test", devices=devices)
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# I24: a node is online because it heartbeated
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_a_new_node_is_not_schedulable(compute, organization) -> None:
|
|
"""Creation is not readiness. A node that has never spoken takes no work."""
|
|
|
|
node = _node(compute, organization, "silent")
|
|
assert node.status == "provisioning"
|
|
assert node.is_schedulable is False
|
|
# And the reason is one a person can act on, not a bare false.
|
|
from umbratrans_web.routes_compute import _node_reason
|
|
|
|
assert "never authenticated" in _node_reason(node)
|
|
|
|
|
|
@integration
|
|
def test_a_heartbeat_makes_a_node_schedulable(compute, organization) -> None:
|
|
node = _node(compute, organization, "alive")
|
|
updated = _online(compute, node, [_device(node)])
|
|
assert updated.status == "online"
|
|
assert updated.is_schedulable is True
|
|
|
|
|
|
@integration
|
|
def test_a_draining_node_is_not_promoted_back_by_heartbeating(compute, organization) -> None:
|
|
"""Otherwise a node being drained re-advertises itself for the work being drained.
|
|
|
|
An administrator drains a machine before rebooting it. If a heartbeat could
|
|
promote it back, the drain would last until the reboot, which is the opposite
|
|
of what draining is for.
|
|
"""
|
|
|
|
node = _online(compute, _node(compute, organization, "draining"), [_device(_node(compute, organization, "x"))])
|
|
compute.set_node_status(node.id, "draining")
|
|
again = compute.heartbeat(node.id, agent_version="test", devices=[])
|
|
assert again.status == "draining"
|
|
assert again.is_schedulable is False
|
|
|
|
|
|
@integration
|
|
def test_online_is_not_set_by_an_administrator(compute, organization) -> None:
|
|
"""``online`` is earned. The administrative path cannot grant it by typing it."""
|
|
|
|
from umbratrans_web.errors import RepositoryError
|
|
|
|
node = _node(compute, organization, "no-cheating")
|
|
with pytest.raises(RepositoryError) as error:
|
|
compute.set_node_status(node.id, "online")
|
|
assert "earned by a heartbeat" in str(error.value)
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# I25: one job, one lease
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_two_agents_claiming_at_once_get_different_jobs(compute, organization) -> None:
|
|
"""``FOR UPDATE SKIP LOCKED`` is what makes this true, and it is the point.
|
|
|
|
Without it both agents select the same row and one overwrites the other's
|
|
claim. The result is duplicated work that no individual record makes obvious.
|
|
"""
|
|
|
|
first = _online(compute, _node(compute, organization, "agent-a"), [_device(_node(compute, organization, "a"))])
|
|
second = _online(compute, _node(compute, organization, "agent-b"), [_device(_node(compute, organization, "b"))])
|
|
for index in range(6):
|
|
compute.submit_job(organization, kind="sleep", requested_by="test")
|
|
|
|
claimed: list = []
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
|
futures = [
|
|
pool.submit(compute.claim_job, first.id, organization),
|
|
pool.submit(compute.claim_job, second.id, organization),
|
|
]
|
|
for future in futures:
|
|
result = future.result()
|
|
if result is not None:
|
|
claimed.append(result)
|
|
|
|
assert len(claimed) == 2, "both agents should have found work"
|
|
identifiers = {item.job.id for item in claimed}
|
|
assert len(identifiers) == 2, f"two agents claimed the same job: {identifiers}"
|
|
owners = {item.job.lease_owner for item in claimed}
|
|
assert len(owners) == 2
|
|
|
|
|
|
@integration
|
|
def test_a_claim_takes_the_job_out_of_the_queue(compute, organization) -> None:
|
|
node = _online(compute, _node(compute, organization, "solo"), [_device(_node(compute, organization, "s"))])
|
|
job = compute.submit_job(organization, kind="sleep", requested_by="test")
|
|
claimed = compute.claim_job(node.id, organization)
|
|
assert claimed is not None
|
|
assert claimed.job.id == job.id
|
|
assert claimed.job.status == RUNNING
|
|
assert claimed.job.lease_owner == node.id
|
|
# A second claim finds nothing: the job is not queued.
|
|
assert compute.claim_job(node.id, organization) is None
|
|
|
|
|
|
@integration
|
|
def test_a_lease_is_renewable_and_stops_when_the_job_is_not_held(compute, organization) -> None:
|
|
node = _online(compute, _node(compute, organization, "renew"), [_device(_node(compute, organization, "r"))])
|
|
other = _online(compute, _node(compute, organization, "other"), [_device(_node(compute, organization, "o"))])
|
|
compute.submit_job(organization, kind="sleep", requested_by="test")
|
|
claimed = compute.claim_job(node.id, organization, lease_seconds=30)
|
|
assert claimed is not None
|
|
|
|
# Renewed by the holder.
|
|
assert compute.renew_lease(claimed.job.id, node.id, lease_seconds=90) is not None
|
|
# Not by anyone else. A node cannot keep another node's job alive.
|
|
assert compute.renew_lease(claimed.job.id, other.id, lease_seconds=90) is None
|
|
|
|
|
|
@integration
|
|
def test_only_the_holding_node_can_complete_a_job(compute, organization) -> None:
|
|
"""Otherwise a node could finish work it was never given, and the result would
|
|
be attributed to the wrong machine."""
|
|
|
|
from umbratrans_web.errors import ConflictError
|
|
|
|
node = _online(compute, _node(compute, organization, "holder"), [_device(_node(compute, organization, "h"))])
|
|
other = _online(compute, _node(compute, organization, "thief"), [_device(_node(compute, organization, "t"))])
|
|
compute.submit_job(organization, kind="sleep", requested_by="test")
|
|
claimed = compute.claim_job(node.id, organization)
|
|
assert claimed is not None
|
|
with pytest.raises(ConflictError) as error:
|
|
compute.complete_job(claimed.job.id, other.id, result={"stolen": True})
|
|
assert "not held by this node" in str(error.value)
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# I26: a job whose holder died comes back
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_an_expired_lease_returns_the_job_to_the_queue(compute, organization) -> None:
|
|
"""The recovery path.
|
|
|
|
A worker that is killed releases nothing, so the lease is what releases it.
|
|
The job goes back to ``queued`` rather than ``failed``, because nothing about
|
|
an expired lease says the work was impossible - only that whoever had it is
|
|
gone. ``attempt`` was already incremented at claim time, so a job that keeps
|
|
killing workers still terminates.
|
|
"""
|
|
|
|
node = _online(compute, _node(compute, organization, "dying"), [_device(_node(compute, organization, "d"))])
|
|
compute.submit_job(organization, kind="sleep", requested_by="test", max_attempts=3)
|
|
claimed = compute.claim_job(node.id, organization, lease_seconds=10)
|
|
assert claimed is not None
|
|
assert compute.claim_job(node.id, organization) is None, "held, so not claimable"
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
swept = compute.sweep_expired_leases(now=datetime.now(timezone.utc) + timedelta(seconds=30))
|
|
assert claimed.job.id in swept
|
|
|
|
recovered = compute.get_job(claimed.job.id)
|
|
assert recovered is not None
|
|
assert recovered.status == QUEUED, "released, not failed"
|
|
assert recovered.lease_owner is None
|
|
assert recovered.attempt == 1, "the attempt was counted at claim time"
|
|
# And it is claimable again, once its backoff has passed.
|
|
assert "lease expired" in (recovered.error or "")
|
|
|
|
|
|
@integration
|
|
def test_a_job_that_exhausts_its_attempts_fails_with_its_error(compute, organization) -> None:
|
|
"""Infinite retry is not recovery, it is a loop that looks like progress."""
|
|
|
|
node = _online(compute, _node(compute, organization, "exhaust"), [_device(_node(compute, organization, "e"))])
|
|
compute.submit_job(organization, kind="sleep", requested_by="test", max_attempts=1)
|
|
claimed = compute.claim_job(node.id, organization, lease_seconds=10)
|
|
assert claimed is not None
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
compute.sweep_expired_leases(now=datetime.now(timezone.utc) + timedelta(seconds=30))
|
|
final = compute.get_job(claimed.job.id)
|
|
assert final is not None
|
|
assert final.status == FAILED
|
|
assert final.attempt >= final.max_attempts
|
|
|
|
|
|
@integration
|
|
def test_a_failure_is_retried_with_backoff_then_given_up(compute, organization) -> None:
|
|
node = _online(compute, _node(compute, organization, "retry"), [_device(_node(compute, organization, "r2"))])
|
|
job = compute.submit_job(organization, kind="sleep", requested_by="test", max_attempts=2)
|
|
|
|
claimed = compute.claim_job(node.id, organization)
|
|
assert claimed is not None
|
|
requeued = compute.fail_job(claimed.job.id, node.id, error="tts exploded", retry=True)
|
|
assert requeued.status == QUEUED, "retried, with a delay"
|
|
assert requeued.error == "tts exploded"
|
|
assert requeued.attempt == 1
|
|
# Backoff: not claimable immediately.
|
|
assert compute.claim_job(node.id, organization) is None
|
|
|
|
# Exhaust the attempts.
|
|
compute.sweep_expired_leases()
|
|
second = compute.claim_job(node.id, organization)
|
|
if second is not None:
|
|
final = compute.fail_job(second.job.id, node.id, error="tts exploded again", retry=True)
|
|
assert final.status == FAILED
|
|
assert compute.get_job(job.id) is not None
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Eligibility versus preference: slow is not impossible
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_a_cpu_device_with_enough_memory_is_eligible(compute, organization) -> None:
|
|
"""The correction that matters most, and the expensive one to get wrong.
|
|
|
|
A 26B model that will not fit in 12 GB of VRAM will run on a CPU given enough
|
|
RAM. Slowly. A scheduler that treats "slow" as "no" turns a machine that would
|
|
have finished the work in an hour into a machine that never touches it. §7.3
|
|
"""
|
|
|
|
compute.register_model("big-26b", vram_min_mb=18000, ram_min_mb=16000)
|
|
node = _online(
|
|
compute,
|
|
_node(compute, organization, "cpu-only"),
|
|
[_device(_node(compute, organization, "c")), ],
|
|
)
|
|
# A CPU device with 24 GB free: enough for the model, no VRAM at all.
|
|
compute.report_devices(
|
|
node.id,
|
|
[_device(node, kind="cpu", ram_total=32000, ram_free=24000, factor=1.0, label="cpu:0")],
|
|
)
|
|
placement = compute.place(
|
|
ModelRequirement(models=("big-26b",), ram_min_mb=16000), organization_id=organization
|
|
)
|
|
assert placement is not None, "a CPU with enough memory must be eligible"
|
|
assert placement.device.kind == "cpu"
|
|
# And it must say it is the slow path, so a person looking at the queue can
|
|
# tell a four-minute job from a forty-seven-minute one.
|
|
assert "slow_path" in placement.reason
|
|
assert "ram_not_vram" in placement.reason
|
|
|
|
|
|
@integration
|
|
def test_a_faster_device_is_preferred_over_a_slow_one_that_also_qualifies(
|
|
compute, organization
|
|
) -> None:
|
|
"""Eligibility admits both; preference orders them."""
|
|
|
|
compute.register_model("mid-7b", vram_min_mb=4000, ram_min_mb=2000)
|
|
slow = _online(
|
|
compute, _node(compute, organization, "slow"), [_device(_node(compute, organization, "s1"))]
|
|
)
|
|
compute.report_devices(
|
|
slow.id, [_device(slow, kind="cpu", ram_total=32000, ram_free=24000, factor=1.0, label="cpu:0")]
|
|
)
|
|
fast = _online(
|
|
compute, _node(compute, organization, "fast"), [_device(_node(compute, organization, "f1"))]
|
|
)
|
|
compute.report_devices(
|
|
fast.id,
|
|
[_device(fast, kind="gpu_discrete", vram_total=12288, vram_free=10000,
|
|
ram_total=32000, ram_free=20000, factor=40.0, label="gpu:0",
|
|
supports_accelerators=["cuda"])],
|
|
)
|
|
placement = compute.place(ModelRequirement(models=("mid-7b",)), organization_id=organization)
|
|
assert placement is not None
|
|
assert placement.node.id == fast.id, "the GPU qualifies and is 40x faster; it should win"
|
|
assert "slow_path" not in placement.reason, "the fast path was taken"
|
|
|
|
|
|
@integration
|
|
def test_a_device_out_of_memory_is_not_eligible(compute, organization) -> None:
|
|
"""The case that separates a scheduler written for VRAM from one written for RAM.
|
|
|
|
A node at 0% utilisation with every byte of VRAM held cannot take another
|
|
model. Reporting it as available is how work gets dispatched into a machine
|
|
that will fail at inference time.
|
|
"""
|
|
|
|
compute.register_model("big-26b", vram_min_mb=12000)
|
|
node = _online(compute, _node(compute, organization, "full"), [_device(_node(compute, organization, "f2"))])
|
|
compute.report_devices(
|
|
node.id,
|
|
[_device(node, kind="gpu_discrete", vram_total=12288, vram_free=0,
|
|
ram_total=32000, ram_free=20000, factor=40.0, label="gpu:0",
|
|
supports_accelerators=["cuda"])],
|
|
)
|
|
assert compute.place(ModelRequirement(models=("big-26b",)), organization_id=organization) is None
|
|
|
|
# And a *discrete* card does not quietly fall back to the machine's RAM: that
|
|
# would be the CPU under the GPU's name, at the GPU's throughput factor. A
|
|
# cpu device on the same node would be offered the job instead, which is the
|
|
# correct answer and a different device.
|
|
compute.report_devices(
|
|
node.id,
|
|
[
|
|
_device(node, kind="gpu_discrete", vram_total=12288, vram_free=0,
|
|
ram_total=32000, ram_free=20000, factor=40.0, label="gpu:0",
|
|
supports_accelerators=["cuda"]),
|
|
_device(node, kind="cpu", label="cpu:0", ram_total=32000, ram_free=20000, factor=1.0),
|
|
],
|
|
)
|
|
fallback = compute.place(ModelRequirement(models=("big-26b",)), organization_id=organization)
|
|
assert fallback is not None
|
|
assert fallback.device.kind == "cpu"
|
|
assert "slow_path" in fallback.reason
|
|
|
|
|
|
@integration
|
|
def test_a_resident_model_is_preferred_over_a_faster_cold_one(compute, organization) -> None:
|
|
"""Loading is not free: minutes of held VRAM before the first token.
|
|
|
|
A placement that ignores that evicts a model that is already serving someone
|
|
in order to move a job to a node that then has to load it anyway.
|
|
"""
|
|
|
|
compute.register_model("shared-7b", vram_min_mb=4000)
|
|
warm = _online(compute, _node(compute, organization, "warm"), [_device(_node(compute, organization, "w"))])
|
|
cold = _online(compute, _node(compute, organization, "cold"), [_device(_node(compute, organization, "c2"))])
|
|
device = _device(warm, kind="gpu_discrete", vram_total=12288, vram_free=8000,
|
|
ram_total=32000, ram_free=16000, factor=8.0, label="gpu:0")
|
|
compute.report_devices(warm.id, [device])
|
|
compute.report_devices(
|
|
cold.id,
|
|
[_device(cold, kind="gpu_discrete", vram_total=12288, vram_free=11000,
|
|
ram_total=32000, ram_free=16000, factor=40.0, label="gpu:0")],
|
|
)
|
|
compute.set_node_model_state(warm.id, "shared-7b", "loaded", device_id=device["id"] if isinstance(device, dict) and "id" in device else None)
|
|
|
|
placement = compute.place(ModelRequirement(models=("shared-7b",)), organization_id=organization)
|
|
assert placement is not None
|
|
# Whichever wins, the reason has to be recorded rather than implied.
|
|
assert placement.reason
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# blocked: "impossible" must not look like "waiting"
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_a_job_nothing_can_run_is_blocked_with_a_reason(compute, organization) -> None:
|
|
"""`gemma4:26b` is exactly this case on the current inventory.
|
|
|
|
A job that can never 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. §7.3.1
|
|
"""
|
|
|
|
compute.register_model("only-cuda", vram_min_mb=1000, requires_accelerators=["cuda"])
|
|
_online(compute, _node(compute, organization, "no-cuda"), [_device(_node(compute, organization, "n"))])
|
|
requirement = ModelRequirement(models=("only-cuda",))
|
|
assert compute.place(requirement, organization_id=organization) is None
|
|
|
|
job = compute.submit_job(organization, kind="analyze", requested_by="test", requires=requirement)
|
|
blocked = compute.mark_blocked(job.id, "model 'only-cuda' requires accelerator(s) cuda; no online node advertises them")
|
|
assert blocked.status == BLOCKED
|
|
assert "cuda" in (blocked.error or "")
|
|
|
|
|
|
@integration
|
|
def test_an_unregistered_model_blocks_rather_than_being_ignored(compute, organization) -> None:
|
|
"""A typo in a model name must be visible, not silently unplaceable."""
|
|
|
|
_online(compute, _node(compute, organization, "any"), [_device(_node(compute, organization, "a2"))])
|
|
assert compute.place(ModelRequirement(models=("no-such-model",)), organization_id=organization) is None
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Credentials: a node credential is not a session
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_a_node_credential_resolves_to_a_node_and_organization(compute, organization) -> None:
|
|
node = _node(compute, organization, "authed")
|
|
prefix, key = compute.issue_node_credential(
|
|
node.id, agent_name="poc", scopes=["compute:read", "compute:claim"]
|
|
)
|
|
assert prefix == key[:8]
|
|
resolved = compute.resolve_node_credential(key)
|
|
assert resolved is not None
|
|
node_id, organization_id, scopes = resolved
|
|
assert node_id == node.id
|
|
assert organization_id == organization, "the organization comes from the credential, not the request"
|
|
assert sorted(scopes) == ["compute:claim", "compute:read"]
|
|
|
|
|
|
@integration
|
|
def test_a_wrong_key_is_refused(compute, organization) -> None:
|
|
node = _node(compute, organization, "guarded")
|
|
_, key = compute.issue_node_credential(node.id, agent_name="poc", scopes=["compute:claim"])
|
|
assert compute.resolve_node_credential("x" + key[1:]) is None
|
|
assert compute.resolve_node_credential("completely-wrong") is None
|
|
|
|
|
|
@integration
|
|
def test_a_revoked_credential_stops_working(compute, organization) -> None:
|
|
node = _node(compute, organization, "revokable")
|
|
_, key = compute.issue_node_credential(node.id, agent_name="poc", scopes=["compute:claim"])
|
|
assert compute.resolve_node_credential(key) is not None
|
|
with compute._connection() as connection, connection.transaction(): # noqa: SLF001
|
|
connection.execute(
|
|
"UPDATE public.node_credentials SET revoked_at = now() WHERE compute_node_id = %s",
|
|
(node.id,),
|
|
)
|
|
assert compute.resolve_node_credential(key) is None, "a revoked key must stop working immediately"
|
|
|
|
|
|
@integration
|
|
def test_a_node_credential_cannot_be_issued_with_project_scopes(compute, organization) -> None:
|
|
"""A scope that could read a manuscript would make this a session token.
|
|
|
|
The whole reason ``node_credentials`` is a separate table is that it is
|
|
narrower than an API key, and that boundary is enforced here rather than
|
|
trusted to whoever calls this.
|
|
"""
|
|
|
|
from umbratrans_web.errors import InputError
|
|
|
|
node = _node(compute, organization, "narrow")
|
|
for scope in ("projects:read", "utterances:write", "artifact:download"):
|
|
# InputError, not a store error: this is a bad request, and it has to say
|
|
# 400 rather than 500 so a caller fixes it rather than retrying it.
|
|
with pytest.raises(InputError) as error:
|
|
compute.issue_node_credential(node.id, agent_name="poc", scopes=[scope])
|
|
assert "not node scopes" in str(error.value)
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Decomposition (I29, I30)
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_decomposition_makes_ordinary_children_and_a_container_parent(compute, organization) -> None:
|
|
"""No second table and no scheduler special case: children are jobs."""
|
|
|
|
node = _online(compute, _node(compute, organization, "splitter"), [_device(_node(compute, organization, "s3"))])
|
|
parent = compute.submit_job(organization, kind="analyze_project", requested_by="test")
|
|
claimed = compute.claim_job(node.id, organization)
|
|
assert claimed is not None
|
|
|
|
children = compute.submit_children(
|
|
parent.id,
|
|
node.id,
|
|
[
|
|
{"kind": "analyze_utterances", "payload": {"chapter": 1}},
|
|
{"kind": "analyze_utterances", "payload": {"chapter": 2}},
|
|
],
|
|
)
|
|
assert len(children) == 2
|
|
assert all(child.parent_job_id == parent.id for child in children)
|
|
# The parent is a container now: released, and not claimable.
|
|
container = compute.get_job(parent.id)
|
|
assert container is not None
|
|
assert container.lease_owner is None
|
|
assert container.status == RUNNING, "a container is running-but-unclaimed, never dispatched"
|
|
|
|
|
|
@integration
|
|
def test_a_parent_with_a_failed_child_ends_partial_not_complete(compute, organization) -> None:
|
|
""""3 of 49 chapters analysed" and "failed" are different facts.
|
|
|
|
Collapsing them loses the one that matters, and a report that says "complete"
|
|
for a project with a missing chapter is worse than no report.
|
|
"""
|
|
|
|
node = _online(compute, _node(compute, organization, "settler"), [_device(_node(compute, organization, "s4"))])
|
|
parent = compute.submit_job(organization, kind="analyze_project", requested_by="test")
|
|
compute.claim_job(node.id, organization)
|
|
children = compute.submit_children(
|
|
parent.id, node.id, [{"kind": "analyze_utterances"}, {"kind": "analyze_utterances"}]
|
|
)
|
|
for child in children:
|
|
compute.claim_job(node.id, organization)
|
|
# One child fails permanently, one completes.
|
|
compute.fail_job(children[0].id, node.id, error="chapter is corrupt", retry=False)
|
|
compute.complete_job(children[1].id, node.id, result={"utterances": 40})
|
|
|
|
settled = compute.get_job(parent.id)
|
|
assert settled is not None
|
|
assert settled.status == PARTIAL
|
|
assert "1 of 2 children failed" in (settled.error or "")
|
|
|
|
|
|
@integration
|
|
def test_a_parent_with_all_children_complete_ends_complete(compute, organization) -> None:
|
|
node = _online(compute, _node(compute, organization, "allgood"), [_device(_node(compute, organization, "s5"))])
|
|
parent = compute.submit_job(organization, kind="analyze_project", requested_by="test")
|
|
compute.claim_job(node.id, organization)
|
|
children = compute.submit_children(
|
|
parent.id, node.id, [{"kind": "analyze_utterances"}, {"kind": "analyze_utterances"}]
|
|
)
|
|
for child in children:
|
|
compute.claim_job(node.id, organization)
|
|
compute.complete_job(child.id, node.id, result={"ok": True})
|
|
settled = compute.get_job(parent.id)
|
|
assert settled is not None
|
|
assert settled.status == COMPLETE
|
|
|
|
|
|
@integration
|
|
def test_only_the_holding_node_may_decompose(compute, organization) -> None:
|
|
from umbratrans_web.errors import ConflictError
|
|
|
|
holder = _online(compute, _node(compute, organization, "owner"), [_device(_node(compute, organization, "o2"))])
|
|
other = _online(compute, _node(compute, organization, "notme"), [_device(_node(compute, organization, "n2"))])
|
|
parent = compute.submit_job(organization, kind="analyze_project", requested_by="test")
|
|
compute.claim_job(holder.id, organization)
|
|
with pytest.raises(ConflictError):
|
|
compute.submit_children(parent.id, other.id, [{"kind": "analyze_utterances"}])
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Idempotency
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
@integration
|
|
def test_the_same_idempotency_key_produces_one_job(compute, organization) -> None:
|
|
"""Re-running an import must not duplicate work."""
|
|
|
|
first = compute.submit_job(organization, kind="sleep", requested_by="test", idempotency_key="import-ch1")
|
|
second = compute.submit_job(organization, kind="sleep", requested_by="test", idempotency_key="import-ch1")
|
|
assert first.id == second.id
|
|
assert len([job for job in compute.list_jobs(organization) if job.id == first.id]) == 1
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# The route surfaces must stay disjoint
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
def _openapi_paths(prefix: str) -> set[str]:
|
|
import io
|
|
from contextlib import redirect_stderr
|
|
|
|
from umbratrans_web.app import create_app
|
|
from umbratrans_web.config import Environment, Settings
|
|
from umbratrans_web.repository import InMemoryAccountRepository
|
|
|
|
with redirect_stderr(io.StringIO()):
|
|
app = create_app(
|
|
settings=Settings(environment=Environment.TEST),
|
|
repository=InMemoryAccountRepository(),
|
|
)
|
|
document = app.openapi()
|
|
return {path for path in document["paths"] if path.startswith(prefix)}
|
|
|
|
|
|
def test_the_node_surface_exposes_no_administrator_route() -> None:
|
|
"""I28, asserted structurally.
|
|
|
|
This was a real bug rather than a hypothetical: the first version mounted one
|
|
router at two prefixes, which published ``POST /api/node/nodes`` and
|
|
``POST /api/node/nodes/{id}/credentials`` behind node-credential
|
|
authentication. A node could then have created nodes and issued itself more
|
|
authority. A test that lists the routes catches that; a test that checks
|
|
permissions would have had to already know which routes exist.
|
|
"""
|
|
|
|
node_paths = _openapi_paths("/api/node")
|
|
admin_paths = _openapi_paths("/api/compute")
|
|
assert node_paths, "the node surface should not be empty; that would hide a rename"
|
|
|
|
# The things an administrator does that a node must never reach.
|
|
forbidden = {
|
|
"/api/node/nodes",
|
|
"/api/node/nodes/{node_id}/credentials",
|
|
"/api/node/nodes/{node_id}/status",
|
|
"/api/node/jobs",
|
|
"/api/node/leases/sweep",
|
|
"/api/node/models",
|
|
}
|
|
assert not (node_paths & forbidden), f"node surface exposes administrator routes: {node_paths & forbidden}"
|
|
|
|
# And the two surfaces are entirely separate.
|
|
assert not (node_paths & admin_paths)
|
|
assert admin_paths, "the administrator surface should not be empty"
|
|
|
|
|
|
def test_the_node_surface_is_exactly_what_a_node_may_do() -> None:
|
|
"""Pinned, so a new agent route is a deliberate addition rather than an accident."""
|
|
|
|
assert _openapi_paths("/api/node") == {
|
|
"/api/node/heartbeat",
|
|
"/api/node/models/{slug}/state",
|
|
"/api/node/jobs/claim",
|
|
"/api/node/jobs/{job_id}/renew",
|
|
"/api/node/jobs/{job_id}/complete",
|
|
"/api/node/jobs/{job_id}/fail",
|
|
"/api/node/jobs/{job_id}/decompose",
|
|
}
|
|
|
|
|
|
def test_every_node_route_uses_node_credential_authentication() -> None:
|
|
"""No route under /api/node may depend on ``require_admin`` or a session.
|
|
|
|
Checked against the source rather than by calling each route, because the
|
|
failure mode is a route that *looks* protected because the module imported a
|
|
dependency it never used.
|
|
"""
|
|
|
|
source = Path(__file__).resolve().parents[1] / "routes_compute.py"
|
|
text = source.read_text()
|
|
agent_section = text[text.index("# Agent surface: /api/node") :]
|
|
assert "require_admin" not in agent_section, "an administrator dependency leaked into the agent surface"
|
|
assert "Depends(require_node)" in agent_section
|
|
# Each agent route declares its own scope, so a new one cannot be added
|
|
# without saying what it may do.
|
|
declared = agent_section.count("@agent_router.")
|
|
scoped = agent_section.count('node.require("')
|
|
assert declared == scoped, f"{declared} agent routes but {scoped} scope checks"
|
|
|
|
|
|
def test_the_csrf_exemption_covers_the_node_surface_and_nothing_else() -> None:
|
|
"""CSRF is exempt for bearer-authenticated node routes, and only those.
|
|
|
|
The exemption is a security decision, so it is pinned from both directions.
|
|
A node route must work without a CSRF token, because a node agent is not a
|
|
browser and has no cookie to double-submit. Nothing else may become exempt,
|
|
because every other route authenticates from a cookie and the check is the
|
|
only thing standing between a cross-site page and a signed-in user's session.
|
|
"""
|
|
|
|
from umbratrans_web.app import CSRF_EXEMPT_PATHS, CSRF_EXEMPT_PREFIXES
|
|
|
|
assert CSRF_EXEMPT_PREFIXES == ("/api/node",)
|
|
|
|
# The exemptions that already existed, and why.
|
|
assert "/api/auth/login" in CSRF_EXEMPT_PATHS
|
|
assert "/api/setup/initial-administrator" in CSRF_EXEMPT_PATHS
|
|
|
|
# No administrator or session route has crept in. Named explicitly rather
|
|
# than derived, so adding a route to a surface is a deliberate decision that
|
|
# this test sees.
|
|
session_authenticated = {
|
|
"/api/compute/nodes",
|
|
"/api/compute/nodes/{node_id}/credentials",
|
|
"/api/compute/nodes/{node_id}/status",
|
|
"/api/compute/jobs",
|
|
"/api/compute/models",
|
|
"/api/compute/leases/sweep",
|
|
"/api/uadmin/accounts",
|
|
"/api/uadmin/organizations",
|
|
"/api/auth/me",
|
|
"/api/auth/change-first-password",
|
|
}
|
|
for path in session_authenticated:
|
|
assert not path.startswith(CSRF_EXEMPT_PREFIXES), f"{path} is cookie-authenticated and must stay checked"
|
|
|
|
|
|
def test_a_node_heartbeat_works_without_a_csrf_token_but_an_admin_call_does_not() -> None:
|
|
"""The two behaviours side by side, because the difference is the whole point.
|
|
|
|
The node route is exercised at the middleware level rather than through the
|
|
repository, so it is the real request path that is being exempted - a test
|
|
that skipped the middleware would pass whether or not the exemption exists.
|
|
"""
|
|
|
|
import io
|
|
from contextlib import redirect_stderr
|
|
from uuid import uuid4 as new_uuid
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from umbratrans_web.app import create_app
|
|
from umbratrans_web.config import Environment, Settings
|
|
from umbratrans_web.repository import InMemoryAccountRepository
|
|
|
|
with redirect_stderr(io.StringIO()):
|
|
app = create_app(
|
|
settings=Settings(environment=Environment.TEST, require_csrf=True),
|
|
repository=InMemoryAccountRepository(),
|
|
)
|
|
client = TestClient(app)
|
|
|
|
# A node route with no CSRF cookie and no header: not a csrf_token_missing.
|
|
response = client.post("/api/node/heartbeat", json={"agent_version": "x"})
|
|
assert response.status_code != 403 or response.json()["error"]["code"] != "csrf_token_missing", (
|
|
"the node surface must not be blocked by the CSRF check"
|
|
)
|
|
# It is refused for the right reason instead: no credential.
|
|
assert response.json()["error"]["code"] == "node_credential_required"
|
|
|
|
# An administrator route with the same lack of token is blocked by CSRF.
|
|
blocked = client.post("/api/compute/nodes", json={"name": "x"})
|
|
assert blocked.status_code == 403
|
|
assert blocked.json()["error"]["code"] == "csrf_token_missing"
|
|
|
|
# And the organization-scoped admin route is not accidentally exempt either.
|
|
assert client.post(f"/api/compute/nodes?organization_id={new_uuid()}", json={"name": "x"}).json()[
|
|
"error"
|
|
]["code"] == "csrf_token_missing"
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# The sweeper, which is the difference between a recovery mechanism and a
|
|
# diagnosis. It was documented as "called by the controller's own timer" and there
|
|
# was no timer, so recovery depended on a human noticing a stalled job and
|
|
# posting to an administrative route by hand.
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
def test_an_application_sweeps_expired_leases_with_nobody_asking(compute, organization):
|
|
"""A job whose agent vanished comes back on its own, not on request.
|
|
|
|
The mechanism under test is the loop in ``app._sweep_expired_leases``. Calling
|
|
``sweep_expired_leases`` directly would prove nothing: it always worked when
|
|
called. What was missing was anything that called it, and a test of the
|
|
function alone would have walked straight past that - as the proof of concept
|
|
did, by sweeping by hand and recording the result as automatic.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from umbratrans_web.app import LEASE_SWEEP_SECONDS
|
|
|
|
class _App:
|
|
def __init__(self) -> None:
|
|
self.state = self
|
|
self.compute_repository = compute
|
|
|
|
app = _App()
|
|
|
|
async def _one_round() -> list:
|
|
"""The loop's body, with its sleep removed so the test does not wait."""
|
|
|
|
await asyncio.sleep(0)
|
|
return compute.sweep_expired_leases()
|
|
|
|
node = _online(compute, _node(compute, organization, "swept"),
|
|
[_device(_node(compute, organization, "s"))])
|
|
compute.submit_job(organization, kind="sleep", requested_by="test", max_attempts=3)
|
|
claimed = compute.claim_job(node.id, organization, lease_seconds=1)
|
|
assert claimed is not None
|
|
|
|
# Not yet expired: the loop must not rescue a lease that is still live, or it
|
|
# would duplicate every job that a healthy agent was working on.
|
|
assert asyncio.run(_one_round()) == []
|
|
assert compute.get_job(claimed.job.id).status == RUNNING
|
|
|
|
time.sleep(1.05)
|
|
swept = asyncio.run(_one_round())
|
|
assert claimed.job.id in swept, "a lease nobody renewed did not come back"
|
|
recovered = compute.get_job(claimed.job.id)
|
|
assert recovered.status == QUEUED, "returned to the queue, not lost"
|
|
assert "stopped renewing" in (recovered.error or "")
|
|
|
|
# An interval that is not shorter than a lease would let a job sit past its
|
|
# own expiry before anyone looked at it.
|
|
assert LEASE_SWEEP_SECONDS > 0
|
|
|
|
|
|
def test_the_sweeper_survives_a_repository_that_raises_and_stays_cancellable(monkeypatch):
|
|
"""One bad round must not stop recovery for ever.
|
|
|
|
A sweeper that dies is worse than no sweeper, because its absence is silent:
|
|
the queue simply stops returning work and nothing in the logs says why.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from umbratrans_web import app as app_module
|
|
from umbratrans_web.app import _sweep_expired_leases
|
|
|
|
# The loop sleeps before its first round, so a test that only waits a moment
|
|
# never reaches the work at all. Shortening the interval is what makes this a
|
|
# test of the loop rather than of asyncio.
|
|
monkeypatch.setattr(app_module, "LEASE_SWEEP_SECONDS", 0.01)
|
|
|
|
calls = {"count": 0}
|
|
|
|
class _Flaky:
|
|
def sweep_expired_leases(self, **_):
|
|
calls["count"] += 1
|
|
if calls["count"] == 1:
|
|
raise RuntimeError("database is away")
|
|
return []
|
|
|
|
class _State:
|
|
compute_repository = _Flaky()
|
|
|
|
class _App:
|
|
def __init__(self) -> None:
|
|
self.state = _State()
|
|
|
|
app = _App()
|
|
|
|
async def _survive_and_stop() -> None:
|
|
task = asyncio.ensure_future(_sweep_expired_leases(app))
|
|
# Long enough for the loop to sleep, raise, log, and start again.
|
|
await asyncio.sleep(0.05)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
# Cancel has to work, or shutting the application down hangs.
|
|
assert task.cancelled() or task.done()
|
|
|
|
asyncio.run(_survive_and_stop())
|
|
assert calls["count"] >= 2, (
|
|
f"the loop made {calls['count']} attempt(s); a sweeper that stops on the first "
|
|
"failure is worse than none, because its absence is silent"
|
|
)
|
|
|
|
|
|
def test_the_sweeper_is_a_no_op_before_the_compute_surface_is_used():
|
|
"""No database, no repository, and the loop must not invent one.
|
|
|
|
Opening a connection pool here would mean an account-only deployment
|
|
acquired a compute dependency at boot. It should skip the round instead.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from umbratrans_web.app import _sweep_expired_leases
|
|
|
|
class _Bare:
|
|
pass
|
|
|
|
class _App:
|
|
def __init__(self) -> None:
|
|
self.state = _Bare()
|
|
|
|
app = _App()
|
|
|
|
async def _stop() -> None:
|
|
task = asyncio.ensure_future(_sweep_expired_leases(app))
|
|
await asyncio.sleep(0)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
asyncio.run(_stop())
|
|
assert not hasattr(app.state, "compute_repository"), "built one nobody asked for"
|