stroblmeandClaude Opus 5
8bd30db016
GPU count is not detected, so FLOW_GPUS was 0 on a fresh install and a node asking for one was silently clamped to zero and ran concurrently with every other. Setting the variable serialised them, but it was an environment variable only — `serve` had --max-runs and --max-workers and no --gpus. The clamp warning now names the flag when nothing here declares a card. The same flags are written into the environment before the settings are built, so a value they refused died in a pydantic import naming no flag. They are checked where they are typed instead. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019Hra4ndWMCLU5F3KjUuVAc
399 lines
13 KiB
Python
399 lines
13 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), device="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), device="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_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
|