Let a node say how much of the machine it takes

Five concurrent training nodes, each sizing its thread pool to every core,
left the engine's own event loop unscheduled: the API stopped answering
within 10 s and every client died. The same shape on a GPU deadlocked a run
for 21 minutes at 0% utilisation with nothing failing and nothing to read --
it just sat in `running`.

@node(resources={"cpus": 2}) is the declaration. The engine holds that much
for the length of the execution, so more of them than the machine has room
for wait their turn rather than oversubscribing it, and a `gpus` node holds
its card exclusively. FLOW_CPUS defaults to every core but two, and those two
are what keeps the engine answering.

Because a thread cap is read when the process imports the library, a warm
worker cannot be told a different one -- so an environment gets a pool of its
own and nodes deriving the same one share it, rather than paying a cold start
per call on exactly the nodes whose imports are slowest. XLA_FLAGS is never
derived: it is a composed, version-dependent string, so it travels in
resources.env where it is visible.

A node that declares nothing is not accounted for and behaves as it always
did -- it just gets FLOW_CPUS/FLOW_MAX_WORKERS as a thread cap, which is the
half of this that fixes the reported incident without anybody declaring
anything. An operator who set OMP_NUM_THREADS themselves still wins.

Resources are claimed strictly before a worker slot, so the two blocking
waits cannot deadlock. A node queued for them publishes node_queued and shows
on GET /workers/resources, because waiting and hanging looked identical.

Accounted, not enforced: no cgroups, no rlimits. Scheduling across machines,
flavours and enforcement are the next steps.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-08-26 21:36:59 +02:00
co-authored by Claude Opus 5
parent 4a38c6ed31
commit 608d30d884
13 changed files with 868 additions and 10 deletions
+249
View File
@@ -0,0 +1,249 @@
"""Who gets the machine: accounting, exclusivity, and what the worker is told.
The failure this exists for is not subtle — five concurrent nodes each sizing a
thread pool to every core starved the engine's own event loop, and three GPU
processes each preallocating most of the card deadlocked at zero utilisation.
Both come down to arithmetic nobody was doing, so the arithmetic is what is
checked here.
"""
import os
import threading
import time
import pytest
from fluksio.flow.resources import (
THREAD_VARS,
Allocation,
ResourceAccountant,
derive_env,
fair_share_env,
)
from fluksio.flow.schemas import NodeDef, Resources
def test_what_is_free_is_what_was_handed_out():
accountant = ResourceAccountant(cpus=4, gpus=0)
with accountant.claim(Resources(cpus=3)):
assert accountant.snapshot()["cpus"] == {"total": 4, "free": 1}
assert accountant.snapshot()["cpus"] == {"total": 4, "free": 4}
def test_a_node_waits_until_there_is_room():
"""Blocking is the mechanism — the same backpressure a worker slot applies."""
accountant = ResourceAccountant(cpus=2)
running = threading.Event()
started = threading.Event()
def second() -> None:
with accountant.claim(Resources(cpus=2), node="study.b"):
started.set()
with accountant.claim(Resources(cpus=2), node="study.a"):
thread = threading.Thread(target=second)
thread.start()
# Long enough to have run if nothing was holding it back.
assert not started.wait(0.2)
assert accountant.snapshot()["waiting"][0]["node"] == "study.b"
running.set()
thread.join(timeout=5)
assert started.is_set()
def test_a_gpu_is_held_by_one_node_at_a_time():
"""The deadlock was three processes each preallocating most of one card."""
accountant = ResourceAccountant(cpus=8, gpus=2)
with accountant.claim(Resources(cpus=1, gpus=1)) as first:
with accountant.claim(Resources(cpus=1, gpus=1)) as second:
assert set(first.gpus) & set(second.gpus) == set()
assert accountant.snapshot()["gpus"] == {"total": 2, "free": 0}
assert accountant.snapshot()["gpus"] == {"total": 2, "free": 2}
def test_asking_for_more_than_there_is_gets_what_there_is():
"""A flow written on a big box still has to run on a laptop."""
accountant = ResourceAccountant(cpus=2, gpus=0)
with accountant.claim(Resources(cpus=64, gpus=4)) as allocation:
assert allocation.cpus == 2
assert allocation.gpus == ()
def test_everything_comes_back_when_a_node_fails():
accountant = ResourceAccountant(cpus=4, gpus=1)
with pytest.raises(ValueError):
with accountant.claim(Resources(cpus=4, gpus=1)):
raise ValueError("the node raised")
assert accountant.snapshot()["cpus"]["free"] == 4
assert accountant.snapshot()["gpus"]["free"] == 1
def test_many_nodes_at_once_all_finish():
"""The accountant must not deadlock under contention; it is on every call."""
accountant = ResourceAccountant(cpus=4)
done = []
def work() -> None:
with accountant.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 accountant.snapshot()["cpus"]["free"] == 4
# -----------------------------------------------------------------------------
# What the worker is started with
# -----------------------------------------------------------------------------
def test_the_share_becomes_the_thread_limit():
env = derive_env(Resources(cpus=3), Allocation(cpus=3))
assert all(env[var] == "3" for var in THREAD_VARS)
assert "CUDA_VISIBLE_DEVICES" not in env
def test_a_gpu_node_is_told_which_card_is_its():
env = derive_env(Resources(cpus=1, gpus=2), Allocation(cpus=1, gpus=(1, 3)))
assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
def test_what_the_node_asked_for_wins():
"""A declaration outranks what the allocation would imply — it is deliberate."""
wanted = Resources(
cpus=4, env={"OMP_NUM_THREADS": "1", "XLA_FLAGS": "--xla_cpu_x=false"}
)
env = derive_env(wanted, Allocation(cpus=4))
assert env["OMP_NUM_THREADS"] == "1"
assert env["XLA_FLAGS"] == "--xla_cpu_x=false"
assert env["MKL_NUM_THREADS"] == "4"
def test_the_shared_pool_divides_what_it_has():
env = fair_share_env(cpus=8, workers=4)
assert all(env[var] == "2" for var in THREAD_VARS)
# Never zero, however many workers there are.
assert fair_share_env(cpus=2, workers=8)["OMP_NUM_THREADS"] == "1"
def test_an_operator_who_set_one_keeps_it(monkeypatch):
"""An explicit value in the engine's environment is an answer, not a default."""
monkeypatch.setitem(os.environ, "OMP_NUM_THREADS", "2")
env = fair_share_env(cpus=16, workers=2)
assert "OMP_NUM_THREADS" not in env
assert env["MKL_NUM_THREADS"] == "8"
# -----------------------------------------------------------------------------
# Declaration
# -----------------------------------------------------------------------------
def test_declaring_nothing_stays_exactly_as_it_was():
"""Every flow that exists today parses unchanged and is not accounted for."""
assert NodeDef(id="poll").resources is None
def test_a_misspelled_resource_is_refused():
with pytest.raises(ValueError):
Resources(cpu=4)
# -----------------------------------------------------------------------------
# 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)
controller = FlowController(store, workers=pool, resources=accountant)
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
+80
View File
@@ -468,3 +468,83 @@ def test_an_emission_on_an_undeclared_port_fails_the_call(pool):
# The worker was retired rather than left mid-generator, so the slot works.
assert run(pool, "def process():\n return {'out': 4}\n") == {"out": 4}
# -----------------------------------------------------------------------------
# Pools per environment
#
# A thread cap is read when the process imports the library, so a warm worker
# cannot be told a different one. A node that declared resources therefore gets
# a pool started with its own environment, and nodes deriving the same one
# share it rather than paying for a cold start each.
# -----------------------------------------------------------------------------
def test_a_declared_environment_gets_a_pool_of_its_own(pool):
child = pool.for_env({"OMP_NUM_THREADS": "2"})
assert child is not pool
assert child.env["OMP_NUM_THREADS"] == "2"
# The same environment is the same pool: warm workers are the point.
assert pool.for_env({"OMP_NUM_THREADS": "2"}) is child
assert pool.for_env({"OMP_NUM_THREADS": "3"}) is not child
# Nothing to derive means the shared pool, with no second process anywhere.
assert pool.for_env({}) is pool
def test_a_child_worker_is_started_with_what_it_was_given(pool):
child = pool.for_env({"OMP_NUM_THREADS": "2"})
seen = child.run(
"demo",
"env",
"import os\n\n\ndef process():\n"
" return {'threads': os.environ.get('OMP_NUM_THREADS', '')}\n",
{},
"demo.env",
timeout=10,
)
assert seen == {"threads": "2"}
def test_retiring_workers_reaches_the_children(pool):
child = pool.for_env({"OMP_NUM_THREADS": "2"})
before = child._generation
pool.respawn_all()
assert child._generation > before
def test_cancelling_reaches_a_node_running_in_a_child(pool):
child = pool.for_env({"OMP_NUM_THREADS": "2"})
started = threading.Event()
failed: list[Exception] = []
def call() -> None:
try:
child.run(
"demo",
"slow",
"import time\n\n\ndef process():\n"
" print('up', flush=True)\n"
" time.sleep(30)\n"
" return {'out': 1}\n",
{},
"demo.slow",
timeout=0,
on_event=lambda _event: started.set(),
)
except Exception as exc: # noqa: BLE001 — the point is that it stopped
failed.append(exc)
thread = threading.Thread(target=call)
thread.start()
# The node has no events, so wait for the process rather than for a frame.
deadline = time.monotonic() + 10
while time.monotonic() < deadline and not child._running:
time.sleep(0.01)
assert pool.cancel("demo.slow"), "the parent must find a child's node"
thread.join(timeout=10)
assert failed