Melvin Strobl
1dc891aa1b
- NodeDef `device`/`device_policy` become `worker`/`worker_policy`, the issue `ignored_device` becomes `ignored_worker`, and the SDK's `@node` takes `worker=` and `worker_policy=` - FlowStore._migrate renames the keys in published and draft documents once, bumping their version; a queued run's `needs` reads the old key - placement passes the label as `label`, client regenerated - docs: worker placement wording, MCP issuer without DOMAIN, device rows in the CLI data directory and the API table, device links in compose comments
419 lines
14 KiB
Python
419 lines
14 KiB
Python
"""Which machine a node runs on, and what happens when none of them is free.
|
|
|
|
The two questions the engine used to answer separately — where does a node run,
|
|
and how much of that machine does it get — are one question here. What is
|
|
checked is that the answer covers every machine attached rather than only the
|
|
host the engine is on, that a claim is given back however the call ends, and
|
|
that a node with nowhere to go waits and says so instead of failing.
|
|
"""
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
from collections.abc import Iterator
|
|
|
|
import pytest
|
|
|
|
from fluksio.flow.placement import Placer
|
|
from fluksio.flow.remote import RemoteWorker, RemoteWorkerHub, WorkerInventory
|
|
from fluksio.flow.resources import ResourceAccountant
|
|
from fluksio.flow.schemas import NodeDef, Resources
|
|
|
|
|
|
@pytest.fixture
|
|
def loop() -> Iterator[asyncio.AbstractEventLoop]:
|
|
"""An event loop running in a thread, as the server's does."""
|
|
running = asyncio.new_event_loop()
|
|
thread = threading.Thread(target=running.run_forever, daemon=True)
|
|
thread.start()
|
|
yield running
|
|
running.call_soon_threadsafe(running.stop)
|
|
thread.join(timeout=5)
|
|
running.close()
|
|
|
|
|
|
class FakeSocket:
|
|
def __init__(self) -> None:
|
|
self.sent: list[dict] = []
|
|
self.arrived = threading.Event()
|
|
|
|
async def send_json(self, payload: dict) -> None:
|
|
self.sent.append(payload)
|
|
self.arrived.set()
|
|
|
|
|
|
def attach(
|
|
hub: RemoteWorkerHub,
|
|
loop: asyncio.AbstractEventLoop,
|
|
name: str = "gpu1",
|
|
labels: tuple[str, ...] = ("gpu",),
|
|
cpus: int = 8,
|
|
gpus: int = 1,
|
|
ram_mb: int | None = None,
|
|
):
|
|
socket = FakeSocket()
|
|
worker = RemoteWorker(
|
|
name=name,
|
|
labels=list(labels),
|
|
send=socket.send_json,
|
|
loop=loop,
|
|
max_parallel=4,
|
|
inventory=WorkerInventory(cpus=cpus, gpus=gpus, ram_mb=ram_mb),
|
|
)
|
|
hub.attach(worker)
|
|
return worker, socket
|
|
|
|
|
|
def placer_over(hub: RemoteWorkerHub | None = None, cpus: int = 4, gpus: int = 0):
|
|
return Placer(local=ResourceAccountant(cpus=cpus, gpus=gpus), hub=hub)
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Where a node goes
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_a_node_asking_for_a_card_finds_the_machine_that_has_one(loop):
|
|
"""No label anywhere — the size of the ask is what places it.
|
|
|
|
This is the whole point of a worker reporting its inventory: the engine has
|
|
no GPU, so a node that needs one used to be clamped down to none and run
|
|
here anyway.
|
|
"""
|
|
hub = RemoteWorkerHub()
|
|
placer = placer_over(hub)
|
|
worker, _ = attach(hub, loop, labels=())
|
|
|
|
with placer.claim(Resources(cpus=2, gpus=1)) as (target, allocation):
|
|
assert target.worker is worker
|
|
assert allocation.gpus == (0,)
|
|
assert worker.accountant.snapshot()["gpus"]["free"] == 0
|
|
|
|
assert worker.accountant.snapshot()["gpus"]["free"] == 1
|
|
|
|
|
|
def test_a_node_that_fits_here_stays_here(loop):
|
|
"""No network, and a worker pool that is already warm."""
|
|
hub = RemoteWorkerHub()
|
|
placer = placer_over(hub)
|
|
attach(hub, loop)
|
|
|
|
with placer.claim(Resources(cpus=2)) as (target, _):
|
|
assert target.worker is None
|
|
assert target.name == "engine"
|
|
|
|
|
|
def test_a_required_label_waits_for_its_machine_rather_than_running_here(loop):
|
|
hub = RemoteWorkerHub()
|
|
placer = placer_over(hub)
|
|
worker, _ = attach(hub, loop, labels=("gpu",))
|
|
|
|
with placer.claim(Resources(cpus=2), label="gpu", policy="require") as (
|
|
target,
|
|
_,
|
|
):
|
|
assert target.worker is worker
|
|
|
|
|
|
def test_a_preferred_label_falls_back_here_and_is_still_accounted(loop):
|
|
"""The fallback used to skip the books entirely, so it held nothing."""
|
|
hub = RemoteWorkerHub()
|
|
placer = placer_over(hub)
|
|
|
|
with placer.claim(Resources(cpus=3), label="gpu", policy="prefer") as (
|
|
target,
|
|
allocation,
|
|
):
|
|
assert target.worker is None
|
|
assert allocation.cpus == 3
|
|
assert placer.local.snapshot()["cpus"]["free"] == 1
|
|
|
|
|
|
def test_asking_for_more_than_anything_has_gets_what_there_is(loop, caplog):
|
|
"""A flow written on a cluster still has to run on a laptop."""
|
|
placer = placer_over(cpus=2)
|
|
|
|
with placer.claim(Resources(cpus=64, gpus=4)) as (target, allocation):
|
|
assert target.name == "engine"
|
|
assert allocation.cpus == 2
|
|
assert allocation.gpus == ()
|
|
|
|
# Cards are declared, not detected, so a machine that has one reads as
|
|
# having none until it is told — and the warning is where that is noticed.
|
|
assert "fluksio serve --gpus" in caplog.text
|
|
|
|
|
|
def test_what_is_clamped_to_is_a_machine_that_exists(loop):
|
|
"""Each dimension taken separately can describe a machine nobody has.
|
|
|
|
A big CPU box and a small GPU box are not one machine with both. Clamping
|
|
to the largest of each would ask for something that never comes free, which
|
|
is the hung run the clamping is there to avoid.
|
|
"""
|
|
hub = RemoteWorkerHub()
|
|
placer = placer_over(hub, cpus=2)
|
|
attach(hub, loop, name="big", labels=(), cpus=64, gpus=0)
|
|
attach(hub, loop, name="small-gpu", labels=(), cpus=4, gpus=1)
|
|
|
|
with placer.claim(Resources(cpus=64, gpus=1)) as (target, allocation):
|
|
assert target.name == "small-gpu"
|
|
assert (allocation.cpus, allocation.gpus) == (4, (0,))
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Waiting
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_a_node_waits_for_a_busy_machine_and_says_so(loop):
|
|
hub = RemoteWorkerHub()
|
|
events: list[dict] = []
|
|
placer = Placer(
|
|
local=ResourceAccountant(cpus=2),
|
|
hub=hub,
|
|
events=type("Bus", (), {"publish": lambda self, e: events.append(e)})(),
|
|
)
|
|
started = threading.Event()
|
|
|
|
def second() -> None:
|
|
with placer.claim(Resources(cpus=2), node="study.b"):
|
|
started.set()
|
|
|
|
with placer.claim(Resources(cpus=2), node="study.a"):
|
|
thread = threading.Thread(target=second, daemon=True)
|
|
thread.start()
|
|
# Long enough to have run if nothing was holding it back.
|
|
assert not started.wait(0.3)
|
|
assert placer.snapshot()["waiting"][0]["node"] == "study.b"
|
|
assert events[0]["type"] == "node_queued"
|
|
assert "2 cpu(s)" in events[0]["detail"]
|
|
|
|
thread.join(timeout=5)
|
|
assert started.is_set()
|
|
assert placer.snapshot()["waiting"] == []
|
|
|
|
|
|
def test_a_reconnecting_client_is_told_what_is_waiting(tmp_path):
|
|
"""The wait is announced once, so the socket's snapshot carries it again.
|
|
|
|
A panel is sent none of it, the same as node status.
|
|
"""
|
|
from fluksio.api.routes.flows import snapshot_payload
|
|
from fluksio.flow.controller import FlowController
|
|
from fluksio.flow.store import FlowStore
|
|
|
|
placer = Placer(local=ResourceAccountant(cpus=2))
|
|
placer._announce("study.b", "r1", time.monotonic(), "waiting for 2 cpu(s)")
|
|
controller = FlowController(FlowStore(tmp_path / "flows"), placer=placer)
|
|
|
|
queued = snapshot_payload(controller)["queued"]
|
|
assert [(e["node"], e["reason"]) for e in queued] == [
|
|
("study.b", "waiting for 2 cpu(s)")
|
|
]
|
|
assert snapshot_payload(controller, set())["queued"] == []
|
|
|
|
|
|
def test_a_worker_attaching_wakes_what_was_waiting_for_one(loop):
|
|
"""A node waiting inside one machine's books could not hear this."""
|
|
hub = RemoteWorkerHub()
|
|
placer = placer_over(hub)
|
|
hub.on_change = placer.wake
|
|
ran = threading.Event()
|
|
|
|
def needs_a_card() -> None:
|
|
with placer.claim(Resources(cpus=1, gpus=1), node="study.fit") as (target, _):
|
|
if target.worker is not None:
|
|
ran.set()
|
|
|
|
thread = threading.Thread(target=needs_a_card, daemon=True)
|
|
thread.start()
|
|
# Nothing has a GPU yet, so it is clamped to none and runs here — the wait
|
|
# only happens while a machine that *could* take it is busy.
|
|
thread.join(timeout=5)
|
|
assert not thread.is_alive()
|
|
|
|
worker, _ = attach(hub, loop, labels=())
|
|
held = worker.accountant.try_take(cpus=1, gpus=1)
|
|
thread = threading.Thread(target=needs_a_card, daemon=True)
|
|
thread.start()
|
|
assert not ran.wait(0.3)
|
|
|
|
worker.accountant.give_back(held)
|
|
thread.join(timeout=5)
|
|
assert ran.is_set()
|
|
|
|
|
|
def test_everything_comes_back_when_a_node_fails(loop):
|
|
placer = placer_over(cpus=4, gpus=1)
|
|
with pytest.raises(ValueError):
|
|
with placer.claim(Resources(cpus=4, gpus=1)):
|
|
raise ValueError("the node raised")
|
|
|
|
assert placer.local.snapshot()["cpus"]["free"] == 4
|
|
assert placer.local.snapshot()["gpus"]["free"] == 1
|
|
|
|
|
|
def test_a_worker_going_away_mid_call_gives_its_share_back(loop):
|
|
hub = RemoteWorkerHub()
|
|
placer = placer_over(hub)
|
|
worker, _ = attach(hub, loop, labels=())
|
|
|
|
with pytest.raises(RuntimeError):
|
|
with placer.claim(Resources(cpus=2, gpus=1)) as (target, _):
|
|
assert target.worker is worker
|
|
hub.detach(worker.name)
|
|
raise RuntimeError("worker went away mid-call")
|
|
|
|
# The books are the worker's own, so they go with it — what matters is that
|
|
# nothing is left holding them and the next claim is not blocked.
|
|
assert worker.accountant.snapshot()["gpus"]["free"] == 1
|
|
|
|
|
|
def test_many_nodes_at_once_all_finish(loop):
|
|
"""The placer must not deadlock under contention; it is on every call."""
|
|
placer = placer_over(cpus=4)
|
|
done = []
|
|
|
|
def work() -> None:
|
|
with placer.claim(Resources(cpus=2), node="study.n"):
|
|
time.sleep(0.01)
|
|
done.append(1)
|
|
|
|
threads = [threading.Thread(target=work) for _ in range(8)]
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join(timeout=10)
|
|
|
|
assert len(done) == 8
|
|
assert placer.local.snapshot()["cpus"]["free"] == 4
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# What a run is told before it starts
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_a_run_is_only_held_for_a_wait_that_will_end(loop):
|
|
"""An ask nothing can grant is cut down and runs; holding it would hang."""
|
|
placer = placer_over(cpus=4)
|
|
|
|
assert placer.satisfiable(None) is None
|
|
assert placer.satisfiable({"cpus": 2}) is None
|
|
# No machine here has a GPU and nothing can be started with one.
|
|
assert placer.satisfiable({"cpus": 1, "gpus": 1}) is None
|
|
|
|
|
|
def test_a_run_waits_for_a_machine_something_can_start(loop):
|
|
class Cluster:
|
|
def covers(self, cpus, gpus, ram_mb, label=None):
|
|
return gpus <= 2
|
|
|
|
def provision(self, cpus, gpus, ram_mb, label=None):
|
|
pass
|
|
|
|
def reconcile(self, attached):
|
|
pass
|
|
|
|
def shapes(self, label=None):
|
|
return [(16, 2, 65536)]
|
|
|
|
def status(self):
|
|
return {"name": "hpc"}
|
|
|
|
placer = placer_over(cpus=4)
|
|
placer.provisioners = [Cluster()]
|
|
|
|
assert placer.satisfiable({"cpus": 1, "gpus": 1}) == "1 gpu(s) and 1 cpu(s)"
|
|
assert placer.satisfiable({"cpus": 1, "gpus": 0}) is None
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# The whole path, through a real worker
|
|
# -----------------------------------------------------------------------------
|
|
|
|
|
|
def test_a_declared_node_is_held_to_its_share(tmp_path):
|
|
"""Four nodes wanting two cores each, on a machine with four.
|
|
|
|
What is checked is the pair: never more than the machine has in flight at
|
|
once, and every one of them told what it was given — the two halves that
|
|
together are the oversubscription this is for. Also that it does not
|
|
deadlock, since the claim is taken before a worker slot and both block.
|
|
"""
|
|
import sys
|
|
|
|
from fluksio.flow.controller import FlowController, RunContext
|
|
from fluksio.flow.messages import DType, MessageSpec
|
|
from fluksio.flow.schemas import FlowDef
|
|
from fluksio.flow.state import MemoryState
|
|
from fluksio.flow.store import FlowStore
|
|
from fluksio.flow.workers import PythonWorkerPool
|
|
|
|
store = FlowStore(tmp_path / "flows")
|
|
store.write_flow(
|
|
FlowDef(
|
|
name="study",
|
|
mode="batch",
|
|
nodes=[
|
|
NodeDef(
|
|
id="fit",
|
|
provides=[MessageSpec(name="threads", dtype=DType.INT)],
|
|
resources=Resources(cpus=2, env={"XLA_FLAGS": "--x=false"}),
|
|
)
|
|
],
|
|
)
|
|
)
|
|
store.write_node_source(
|
|
"study",
|
|
"fit",
|
|
"import os, time\n\n\ndef process():\n"
|
|
" time.sleep(0.2)\n"
|
|
" return {'threads': int(os.environ['OMP_NUM_THREADS'])}\n",
|
|
)
|
|
|
|
pool = PythonWorkerPool(python=sys.executable, size=4)
|
|
pool.start()
|
|
accountant = ResourceAccountant(cpus=4)
|
|
placer = Placer(local=accountant)
|
|
controller = FlowController(
|
|
store, workers=pool, resources=accountant, placer=placer
|
|
)
|
|
answers: list[object] = []
|
|
in_flight: list[int] = []
|
|
|
|
def once(index: int) -> None:
|
|
pipeline = controller.build_run_pipeline(
|
|
store.read_flow("study"),
|
|
state=MemoryState(),
|
|
run=RunContext(run_id=f"r{index}"),
|
|
)
|
|
pipeline.run({})
|
|
answers.append(pipeline.state.get("study.threads"))
|
|
|
|
def watch(until: threading.Event) -> None:
|
|
while not until.is_set():
|
|
in_flight.append(4 - int(accountant.snapshot()["cpus"]["free"]))
|
|
time.sleep(0.01)
|
|
|
|
finished = threading.Event()
|
|
watcher = threading.Thread(target=watch, args=(finished,))
|
|
watcher.start()
|
|
threads = [threading.Thread(target=once, args=(i,)) for i in range(4)]
|
|
try:
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join(timeout=60)
|
|
assert not thread.is_alive(), "a claim and a worker slot deadlocked"
|
|
finally:
|
|
finished.set()
|
|
watcher.join(timeout=5)
|
|
pool.stop()
|
|
|
|
assert answers == [2, 2, 2, 2]
|
|
assert max(in_flight) <= 4
|
|
assert accountant.snapshot()["cpus"]["free"] == 4
|
|
# One environment, so one extra pool however many nodes derived it.
|
|
assert len(pool._children) == 1
|