396 lines
12 KiB
Python
396 lines
12 KiB
Python
"""The GPU target: one user-chosen destination, resolved honestly.
|
|
|
|
The property that matters most: `decide()` is the single answer used by both
|
|
the generation path and the header badge. If they could disagree, the badge
|
|
would claim work goes somewhere it does not — worse than showing nothing.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from worker import registry, routing
|
|
from worker.errors import ErrorClass, WorkerError
|
|
from worker.identity import WorkerKeypair, issue_session
|
|
from worker.pool import WorkerPool
|
|
|
|
|
|
@pytest.fixture
|
|
def db(tmp_path, monkeypatch):
|
|
from worker import registry as reg
|
|
|
|
db_globals = reg.db_conn.__wrapped__.__globals__
|
|
path = str(tmp_path / "userdata.db")
|
|
with sqlite3.connect(path) as conn:
|
|
conn.executescript(db_globals["_BASE_SCHEMA"])
|
|
monkeypatch.setitem(db_globals, "DB_PATH", path)
|
|
return path
|
|
|
|
|
|
@pytest.fixture
|
|
def settings(monkeypatch):
|
|
"""In-memory settings, so a target choice does not touch the real store."""
|
|
store: dict[str, str] = {}
|
|
monkeypatch.setattr("services.settings_store.get_text", lambda k, d=None: store.get(k, d))
|
|
monkeypatch.setattr("services.settings_store.set_text", lambda k, v: store.__setitem__(k, v))
|
|
return store
|
|
|
|
|
|
class _Plane:
|
|
"""A control plane stub with a real pool, so availability is real logic."""
|
|
|
|
def __init__(self, running=True):
|
|
self.running = running
|
|
self.pool = WorkerPool()
|
|
|
|
|
|
def _enroll(name="desktop-4090", **kw):
|
|
return registry.enroll_worker(
|
|
name=name, public_key=WorkerKeypair.generate().public_bytes(), **kw
|
|
)
|
|
|
|
|
|
def _connect(plane, record, *, now=None):
|
|
# Real clock: ConnectedWorker.stale() compares against now(), so a pinned
|
|
# 1970 timestamp would make every worker look dead.
|
|
now = time.time() if now is None else now
|
|
return plane.pool.connect(
|
|
record,
|
|
session=issue_session(worker_id=record.id, key_id=record.key_id, epoch=1, now=now),
|
|
epoch=1,
|
|
backend="cuda",
|
|
now=now,
|
|
)
|
|
|
|
|
|
# ── The choice ─────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_defaults_to_local(settings):
|
|
assert routing.get_target_id() == routing.LOCAL
|
|
|
|
|
|
def test_choice_is_persisted(settings):
|
|
"""A target that resets on restart would quietly send work elsewhere."""
|
|
routing.set_target_id("abc123")
|
|
assert routing.get_target_id() == "abc123"
|
|
assert settings["worker_target"] == "abc123"
|
|
|
|
|
|
def test_empty_choice_means_local(settings):
|
|
assert routing.set_target_id("") == routing.LOCAL
|
|
assert routing.get_target_id() == routing.LOCAL
|
|
|
|
|
|
def test_a_broken_settings_store_falls_back_to_local(monkeypatch):
|
|
def _boom(*a, **k):
|
|
raise RuntimeError("no db")
|
|
|
|
monkeypatch.setattr("services.settings_store.get_text", _boom)
|
|
assert routing.get_target_id() == routing.LOCAL
|
|
|
|
|
|
# ── The list ───────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_local_is_always_first_and_always_available(db, settings):
|
|
targets = routing.list_targets(_Plane(running=False))
|
|
assert targets[0].id == routing.LOCAL
|
|
assert targets[0].available is True
|
|
|
|
|
|
def test_enrolled_workers_are_listed(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
_connect(plane, worker)
|
|
|
|
targets = routing.list_targets(plane)
|
|
|
|
assert [t.id for t in targets] == [routing.LOCAL, worker.id]
|
|
assert targets[1].label == "desktop-4090"
|
|
assert targets[1].available is True
|
|
|
|
|
|
def test_connected_worker_target_exposes_heartbeat_telemetry(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll(
|
|
"desktop-4090",
|
|
host={
|
|
"cpu_count": 32,
|
|
"system_memory_bytes": 64 * 1024**3,
|
|
"gpus": [{"model": "NVIDIA RTX 4090", "memory_bytes": 24 * 1024**3}],
|
|
},
|
|
)
|
|
_connect(plane, worker)
|
|
plane.pool.heartbeat(
|
|
worker.id,
|
|
active_tasks=0,
|
|
available_slots=1,
|
|
free_memory_bytes=20 * 1024**3,
|
|
cpu_percent=37.5,
|
|
)
|
|
|
|
target = routing.list_targets(plane)[1].to_dict()
|
|
|
|
assert target["cpu_percent"] == 37.5
|
|
assert target["free_memory_bytes"] == 20 * 1024**3
|
|
assert target["gpu_name"] == "NVIDIA RTX 4090"
|
|
assert target["gpu_memory_bytes"] == 24 * 1024**3
|
|
|
|
|
|
def test_target_distinguishes_unknown_vram_from_a_full_gpu(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll(
|
|
"desktop-4090",
|
|
host={"gpus": [{"model": "NVIDIA RTX 4090", "memory_bytes": 24 * 1024**3}]},
|
|
)
|
|
_connect(plane, worker)
|
|
|
|
unknown = routing.list_targets(plane)[1].to_dict()
|
|
assert unknown["free_memory_bytes"] is None
|
|
|
|
plane.pool.heartbeat(
|
|
worker.id,
|
|
active_tasks=0,
|
|
available_slots=1,
|
|
free_memory_bytes=0,
|
|
)
|
|
full = routing.list_targets(plane)[1].to_dict()
|
|
assert full["free_memory_bytes"] == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"setup,detail",
|
|
[
|
|
(lambda w: registry.set_enabled(w.id, False), "disabled"),
|
|
(lambda w: registry.revoke(w.id), "removed"),
|
|
],
|
|
)
|
|
def test_unusable_workers_say_why(db, settings, setup, detail):
|
|
"""A greyed entry has to name something the user can act on."""
|
|
plane = _Plane()
|
|
worker = _enroll()
|
|
_connect(plane, worker)
|
|
setup(worker)
|
|
|
|
listed = {t.id: t for t in routing.list_targets(plane)}
|
|
entry = listed.get(worker.id)
|
|
if detail == "removed":
|
|
# Revoked workers drop out of the list entirely.
|
|
assert entry is None
|
|
return
|
|
assert entry.available is False
|
|
assert entry.detail == detail
|
|
|
|
|
|
def test_a_disconnected_worker_is_listed_but_offline(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll()
|
|
|
|
entry = [t for t in routing.list_targets(plane) if t.id == worker.id][0]
|
|
|
|
assert entry.connected is False
|
|
assert entry.available is False
|
|
assert entry.detail == "offline"
|
|
|
|
|
|
def test_an_unapproved_worker_is_not_available(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll(consent_granted=False)
|
|
_connect(plane, worker)
|
|
|
|
entry = [t for t in routing.list_targets(plane) if t.id == worker.id][0]
|
|
assert entry.available is False
|
|
assert entry.detail == "not approved"
|
|
|
|
|
|
def test_a_paused_worker_reports_its_breaker(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll()
|
|
_connect(plane, worker)
|
|
# Real clock again: a breaker opened at a 1970 timestamp would have its
|
|
# cooldown long expired by now and would read as closed.
|
|
for _ in range(3):
|
|
plane.pool.breakers.record_failure(
|
|
worker.id,
|
|
"e:m",
|
|
WorkerError(error_class=ErrorClass.TRANSIENT, code="X", message="x"),
|
|
now=time.time(),
|
|
)
|
|
|
|
entry = [t for t in routing.list_targets(plane) if t.id == worker.id][0]
|
|
assert entry.available is False
|
|
assert "paused" in entry.detail
|
|
|
|
|
|
# ── The decision ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_local_choice_runs_locally(db, settings):
|
|
assert routing.decide(_Plane()).remote is False
|
|
|
|
|
|
def test_chosen_and_reachable_runs_remotely(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
_connect(plane, worker)
|
|
routing.set_target_id(worker.id)
|
|
|
|
decision = routing.decide(plane)
|
|
|
|
assert decision.remote is True
|
|
assert decision.worker_id == worker.id
|
|
assert decision.label == "desktop-4090"
|
|
|
|
|
|
def test_only_the_chosen_worker_is_used(db, settings):
|
|
"""Others may be connected; standby means they receive nothing."""
|
|
plane = _Plane()
|
|
chosen = _enroll("chosen")
|
|
standby = _enroll("standby")
|
|
_connect(plane, chosen)
|
|
_connect(plane, standby)
|
|
routing.set_target_id(chosen.id)
|
|
|
|
assert routing.decide(plane).worker_id == chosen.id
|
|
|
|
|
|
def test_an_offline_choice_falls_back_locally_and_says_so(db, settings):
|
|
"""A user whose remote GPU went to sleep should get their audio, not an
|
|
error about infrastructure."""
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
routing.set_target_id(worker.id)
|
|
|
|
decision = routing.decide(plane)
|
|
|
|
assert decision.remote is False
|
|
assert "desktop-4090" in decision.reason
|
|
assert "offline" in decision.reason
|
|
|
|
|
|
def test_feature_off_falls_back_locally(db, settings):
|
|
worker = _enroll()
|
|
routing.set_target_id(worker.id)
|
|
decision = routing.decide(_Plane(running=False))
|
|
assert decision.remote is False
|
|
assert "turned off" in decision.reason
|
|
|
|
|
|
def test_a_deleted_choice_falls_back_rather_than_stranding(db, settings):
|
|
routing.set_target_id("nosuchworker")
|
|
decision = routing.decide(_Plane())
|
|
assert decision.remote is False
|
|
assert "no longer exists" in decision.reason
|
|
|
|
|
|
def test_a_revoked_choice_falls_back(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll()
|
|
_connect(plane, worker)
|
|
routing.set_target_id(worker.id)
|
|
registry.revoke(worker.id)
|
|
|
|
assert routing.decide(plane).remote is False
|
|
|
|
|
|
# ── Per-operation coverage ─────────────────────────────────────────────────
|
|
|
|
|
|
def test_a_ported_operation_still_runs_remotely(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
_connect(plane, worker)
|
|
routing.set_target_id(worker.id)
|
|
|
|
assert routing.decide(plane, op="tts").remote is True
|
|
|
|
|
|
def test_dubbing_claims_the_remote_target_after_its_producer_lands(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
_connect(plane, worker)
|
|
routing.set_target_id(worker.id)
|
|
|
|
decision = routing.decide(plane, op="dub")
|
|
|
|
assert decision.remote is True
|
|
assert decision.worker_id == worker.id
|
|
|
|
|
|
def test_dubbing_reports_reachability_after_its_producer_lands(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
routing.set_target_id(worker.id)
|
|
|
|
assert "offline" in routing.decide(plane, op="dub").reason
|
|
|
|
|
|
def test_omitting_the_operation_answers_for_the_target(db, settings):
|
|
"""The picker's own menu asks about the machine, not about one job."""
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
_connect(plane, worker)
|
|
routing.set_target_id(worker.id)
|
|
|
|
assert routing.decide(plane).remote is True
|
|
assert routing.supports_operation(None) is True
|
|
|
|
|
|
def test_unknown_operations_are_not_remote(db, settings):
|
|
assert routing.supports_operation("tts") is True
|
|
assert routing.supports_operation("longform") is True
|
|
assert routing.supports_operation("batch") is True
|
|
assert routing.supports_operation("dictation") is False
|
|
assert routing.supports_operation("something-new") is False
|
|
|
|
|
|
def test_status_answers_for_the_operation_it_was_asked_about(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
_connect(plane, worker)
|
|
routing.set_target_id(worker.id)
|
|
|
|
payload = routing.status(plane, op="dub")
|
|
|
|
assert payload["op"] == "dub"
|
|
assert payload["target"] == worker.id, "the user's choice is not lost"
|
|
assert payload["active"]["remote"] is True
|
|
# What the menu needs to say "gpu2 · TTS only".
|
|
assert payload["remote_operations"] == [
|
|
"audiobook", "batch", "batch_segments", "clone", "dub", "dub_segments",
|
|
"longform", "tts",
|
|
]
|
|
|
|
|
|
# ── The badge cannot lie ───────────────────────────────────────────────────
|
|
|
|
|
|
def test_status_agrees_with_decide(db, settings):
|
|
"""One answer, used by both the router and the badge."""
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
_connect(plane, worker)
|
|
routing.set_target_id(worker.id)
|
|
|
|
payload = routing.status(plane)
|
|
|
|
assert payload["target"] == worker.id
|
|
assert payload["active"] == routing.decide(plane).to_dict()
|
|
assert payload["active"]["remote"] is True
|
|
|
|
|
|
def test_status_shows_local_when_the_choice_is_unreachable(db, settings):
|
|
plane = _Plane()
|
|
worker = _enroll("desktop-4090")
|
|
routing.set_target_id(worker.id)
|
|
|
|
payload = routing.status(plane)
|
|
|
|
# The user's CHOICE is still the worker — that is not lost — but the
|
|
# ACTIVE answer is local, and the badge must show the active one.
|
|
assert payload["target"] == worker.id
|
|
assert payload["active"]["remote"] is False
|