Fluksio/app
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Let a node say how much of the machine it takes

Browse Source 
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:
stroblme
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
Download Patch File Download Diff File Expand all files Collapse all files
backend/fluksio/flow/controller.py
+67
View File
@@ -62,6 +62,7 @@ from fluksio.flow.pipeline import (
ValueSource,
)
from fluksio.flow.remote import RemoteWorkerHub
from fluksio.flow.resources import ResourceAccountant, derive_env
from fluksio.flow.schemas import (
BrainEdge,
BrainGraph,
@@ -71,6 +72,7 @@ from fluksio.flow.schemas import (
NodeDef,
NodeStatusPublic,
NodeTypeInfo,
Resources,
)
from fluksio.flow.secrets import SecretNotFound, resolve_params
from fluksio.flow.state import MemoryState, StateBackend
@@ -341,6 +343,7 @@ class FlowController:
alerts: AlertManager | None = None,
workers: PythonWorkerPool | None = None,
remote: RemoteWorkerHub | None = None,
resources: ResourceAccountant | None = None,
) -> None:
self.store = store
# Without a pool, python nodes are compiled and run in this process —
@@ -348,6 +351,9 @@ class FlowController:
self.workers = workers
# Workers on other hosts. A node without a device never touches it.
self.remote = remote
# What the machine has, for the nodes that say what they need. Without
# one, a declaration is recorded and nothing is held against it.
self.resources = resources
self.state = state if state is not None else MemoryState()
self.events = events
self.max_workers = max_workers
@@ -790,6 +796,56 @@ class FlowController:
# Building
# -------------------------------------------------------------------------
def _allocated(
self,
wanted: Resources,
owner: str,
local: str,
code: str,
*,
node_id: str,
timeout: float,
run_id: str,
on_event: Callable[[dict[str, Any]], None],
) -> Callable[..., Any]:
"""A call that holds its share of the machine while it runs.
The order is load-bearing: the resources are claimed first, and only
then is a worker slot taken. The other way round, a node holding a slot
could sit waiting for cores that a node holding the cores cannot get a
slot to release.
The worker comes from the pool whose environment this allocation
derives, so what the node is told about its share is what the library
inside it reads at import — the only moment those variables are read.
"""
if self.workers is None or self.resources is None:
return self.workers.proxy( # type: ignore[union-attr]
owner,
local,
code,
node_id=node_id,
timeout=timeout,
run_id=run_id,
on_event=on_event,
)
accountant, pool = self.resources, self.workers
def call(**kwargs: Any) -> Any:
with accountant.claim(wanted, node=node_id, run=run_id) as allocation:
return pool.for_env(derive_env(wanted, allocation)).run(
owner,
local,
code,
kwargs,
node_id,
timeout,
run_id=run_id,
on_event=on_event,
)
return call
def _build_flows(
self,
flows: list[tuple[FlowDef, bool]],
@@ -894,6 +950,17 @@ class FlowController:
run_id=run.run_id if run else "",
on_event=emissions.handle,
)
if node_def.resources is not None and not node_def.device:
function = self._allocated(
node_def.resources,
owner,
local,
code,
node_id=node_id,
timeout=timeout,
run_id=run.run_id if run else "",
on_event=emissions.handle,
)
if node_def.device and self.remote is not None:
# A node with a device runs on a worker carrying that
# label. Which worker is decided per call, so one that
backend/fluksio/flow/resources.py
+244
View File
@@ -0,0 +1,244 @@
"""Who gets the machine, and what a node is told about the share it got.
Two halves of one problem. The engine's own event loop has to keep answering
while nodes run, and a numerical library left to itself sizes its thread pool
to every core on the box — so five concurrent nodes were five processes each
believing they owned the machine, and the API stopped answering inside ten
seconds. On a GPU the same shape is worse: three processes each preallocating
most of the card deadlock at zero utilisation, with nothing failing and nothing
to read.
:class:`ResourceAccountant` is the first half: a node that declares what it
needs waits until that much is free, the way it already waits for a worker
slot. :func:`derive_env` is the second: what it got is handed to the library as
the environment it reads at import, because that is the only moment those knobs
are read.
Deliberately cooperative — nothing here is enforced with cgroups or rlimits, so
a node that ignores its share is only accounted for, not stopped. That is the
same trust the worker pool already extends to node code, and the enforcement
half belongs with the scheduler work this is the first step of.
"""
from __future__ import annotations
import logging
import os
import threading
import time
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from fluksio.flow.events import EventBus
from fluksio.flow.schemas import Resources
logger = logging.getLogger(__name__)
#: Every spelling of "how many threads may you use" that a scientific stack
#: reads out of the environment at import. Set together, because a process
#: usually pulls in more than one of them.
THREAD_VARS = (
"OMP_NUM_THREADS",
"OPENBLAS_NUM_THREADS",
"MKL_NUM_THREADS",
"VECLIB_MAXIMUM_THREADS",
"NUMEXPR_NUM_THREADS",
)
#: Cores left for the engine when the inventory is worked out rather than
#: configured. The event loop, the queue, the API and the run drivers all live
#: in that process, and a node is not allowed to take the last of it.
ENGINE_RESERVE = 2
def machine_cpus() -> int:
"""Cores a node may be given here, when nobody said."""
return max(1, (os.cpu_count() or 1) - ENGINE_RESERVE)
@dataclass(frozen=True, slots=True)
class Allocation:
"""What one execution was actually given."""
cpus: int = 1
gpus: tuple[int, ...] = ()
@dataclass
class _Waiting:
node: str
since: float
reason: str
class ResourceAccountant:
"""What is free on this machine, and who is waiting for it.
A plain condition variable over two counters. Blocking is the whole
mechanism: the caller is a node thread that would otherwise be executing,
and making it wait is the same backpressure that the worker pool's slot
queue already applies — one reason the two must always be taken in the same
order, resources first, so neither can be held while waiting for the other.
"""
def __init__(
self, cpus: int = 0, gpus: int = 0, events: EventBus | None = None
) -> None:
self.cpus = cpus or machine_cpus()
self.gpus = max(0, gpus)
self.events = events
self._free_cpus = self.cpus
self._free_gpus = list(range(self.gpus))
self._waiting: dict[int, _Waiting] = {}
self._ticket = 0
self._condition = threading.Condition()
# -- what a caller asks for ------------------------------------------------
def _clamp(self, wanted: Resources, node: str) -> tuple[int, int]:
"""What can be granted here, which may be less than what was asked.
A flow written against a sixty-four core box should still run on a
laptop. Waiting forever for cores that do not exist is not a smaller
machine, it is a hung run.
"""
cpus = min(wanted.cpus, self.cpus)
gpus = min(wanted.gpus, self.gpus)
if cpus != wanted.cpus or gpus != wanted.gpus:
logger.warning(
"%s asked for %d cpu(s) and %d gpu(s); this engine has %d and %d",
node or "a node",
wanted.cpus,
wanted.gpus,
self.cpus,
self.gpus,
)
return cpus, gpus
@contextmanager
def claim(
self, wanted: Resources, node: str = "", run: str = ""
) -> Iterator[Allocation]:
"""Hold this node's share for as long as it runs."""
cpus, gpus = self._clamp(wanted, node)
allocation = self._take(cpus, gpus, node, run)
try:
yield allocation
finally:
self._give_back(allocation)
def _take(self, cpus: int, gpus: int, node: str, run: str) -> Allocation:
waited_from = time.monotonic()
ticket = 0
with self._condition:
while self._free_cpus < cpus or len(self._free_gpus) < gpus:
if not ticket:
reason = (
f"waiting for {cpus} cpu(s) ({self._free_cpus} free)"
if self._free_cpus < cpus
else f"waiting for {gpus} gpu(s) ({len(self._free_gpus)} free)"
)
ticket = self._announce(node, run, waited_from, reason)
self._condition.wait()
self._free_cpus -= cpus
taken = tuple(self._free_gpus.pop(0) for _ in range(gpus))
if ticket:
self._waiting.pop(ticket, None)
if ticket:
logger.info(
"%s waited %.1fs for %d cpu(s) and %d gpu(s)",
node or "a node",
time.monotonic() - waited_from,
cpus,
gpus,
)
return Allocation(cpus=cpus, gpus=taken)
def _give_back(self, allocation: Allocation) -> None:
with self._condition:
self._free_cpus += allocation.cpus
self._free_gpus.extend(allocation.gpus)
self._free_gpus.sort()
self._condition.notify_all()
# -- what it looks like from outside ---------------------------------------
def _announce(self, node: str, run: str, since: float, reason: str) -> int:
"""Say a node is queued, not stuck.
The failure this exists for looked identical to a hang: a run sitting
at `running` for twenty minutes with no error and no output. A node
waiting its turn has to say so somewhere a person will look.
"""
self._ticket += 1
ticket = self._ticket
self._waiting[ticket] = _Waiting(node=node, since=since, reason=reason)
logger.info("%s is %s", node or "a node", reason)
if self.events is not None:
self.events.publish(
{
"type": "node_queued",
"flow": node.split(".", 1)[0] if node else "",
"node": node,
"run": run,
"detail": reason,
"ts": time.time(),
}
)
return ticket
def snapshot(self) -> dict[str, object]:
"""What is free and who is waiting, for the workers screen."""
with self._condition:
free_cpus, free_gpus = self._free_cpus, len(self._free_gpus)
return {
"cpus": {"total": self.cpus, "free": free_cpus},
"gpus": {"total": self.gpus, "free": free_gpus},
"waiting": [
{
"node": entry.node,
"reason": entry.reason,
"seconds": round(time.monotonic() - entry.since, 1),
}
for entry in list(self._waiting.values())
],
}
# -----------------------------------------------------------------------------
# What the worker is told
# -----------------------------------------------------------------------------
def derive_env(wanted: Resources, allocation: Allocation) -> dict[str, str]:
"""The environment a worker running this node is started with.
Precedence is engine environment, then what the allocation implies, then
what the node asked for — a declaration is a deliberate statement about
this node and outranks the machine's own default.
``XLA_FLAGS`` is deliberately not derived. It is one composed string whose
contents depend on the version installed, so writing it here would silently
replace whatever the author had put there. It travels in ``resources.env``,
where it is visible.
"""
env = {var: str(allocation.cpus) for var in THREAD_VARS}
if wanted.gpus:
env["CUDA_VISIBLE_DEVICES"] = ",".join(str(index) for index in allocation.gpus)
env.update(wanted.env)
return env
def fair_share_env(cpus: int, workers: int) -> dict[str, str]:
"""Thread caps for the shared pool, where nothing declared anything.
The pool can have every one of its workers busy at once, so each of them
believing it owns the machine is the oversubscription this whole module is
about — and a node that declared nothing is the common case, not the
exception. Only vars the operator has not already set: an explicit
``OMP_NUM_THREADS`` in the engine's environment is an answer, and this is
a default.
"""
fair = max(1, cpus // max(1, workers))
return {var: str(fair) for var in THREAD_VARS if var not in os.environ}
backend/fluksio/flow/schemas.py
+56 -1
View File
@@ -9,7 +9,7 @@ from __future__ import annotations
import re
from typing import Any, Literal
from pydantic import BaseModel, Field, field_validator
from pydantic import BaseModel, ConfigDict, Field, field_validator
from fluksio.flow.messages import MessageSpec
@@ -24,6 +24,53 @@ def _validate_name(value: str) -> str:
return value
class Resources(BaseModel):
"""What one execution of a node needs to have to itself.
Declaring nothing is the default and means what it always did: the node
runs on the shared worker pool and nothing is accounted for it. That is
right for the kind of node most flows are made of — a poll, a threshold, a
message on its way somewhere.
It is wrong for the other kind. A numerical library sizes its thread pool
to every core it can see, so a handful of them at once oversubscribe the
machine badly enough to starve the engine's own event loop, and a GPU
library that preallocates most of the card deadlocks when a second one
arrives. Both are a node saying how much of the machine it takes, which is
what this is.
"""
model_config = ConfigDict(extra="forbid")
cpus: int = Field(
default=1,
ge=1,
description=(
"Cores held for the whole execution. Also what the thread-pool "
"variables are set to, so a library sizing itself to the machine "
"sizes itself to this instead."
),
)
gpus: int = Field(
default=0,
ge=0,
description=(
"Whole devices held for the whole execution, named to the node "
"through CUDA_VISIBLE_DEVICES. Nothing else is given them while it "
"runs, which is what keeps two preallocating processes apart."
),
)
env: dict[str, str] = Field(
default_factory=dict,
description=(
"Extra environment for the worker this node runs in, applied over "
"what the allocation derives. Where a library's own tuning goes — "
"XLA_FLAGS, XLA_PYTHON_CLIENT_MEM_FRACTION — since those are "
"composed strings the engine must not invent."
),
)
class NodeDef(BaseModel):
"""A node as stored: identity, configuration and ports.
@@ -78,6 +125,14 @@ class NodeDef(BaseModel):
"function whose answer can change on its own."
),
)
resources: Resources | None = Field(
default=None,
description=(
"What one execution of this node holds while it runs. Absent — the "
"default — means it is not accounted for and shares the engine's "
"workers, which is right for everything that is not compute-heavy."
),
)
@field_validator("id")
@classmethod
backend/fluksio/flow/workers.py
+49 -6
View File
@@ -228,6 +228,10 @@ class PythonWorkerPool:
# Put into every worker's environment past the denylist — where the
# artifact store is, which node code needs and cannot guess.
self.env = env or {}
# Pools for nodes that declared resources, keyed by the environment
# their allocation derived. Empty on all but the one pool the engine
# builds; see ``for_env``.
self._children: dict[frozenset[tuple[str, str]], PythonWorkerPool] = {}
self._idle: queue.Queue[_Worker | None] = queue.Queue()
# Keyed by (run, node): a sweep has the same node executing in several
# runs at once, and cancelling one of them must not kill the others.
@@ -255,9 +259,41 @@ class PythonWorkerPool:
for _ in range(self.size):
self._idle.put(None)
def for_env(self, extra: dict[str, str]) -> PythonWorkerPool:
"""The pool whose workers were started with this environment.
A thread cap or a GPU assignment is read once, when the process
imports the library, so a worker that is already warm cannot be told a
different one — and retiring a worker per call would give back the cold
starts ``_warm`` exists to avoid, on exactly the nodes whose imports
are slowest. So an environment gets a pool of its own instead, and
nodes deriving the same one share it.
How many of them may run at once is the accountant's business, not this
pool's: the slots here only bound one environment's own concurrency.
"""
if not extra:
return self
key = frozenset(extra.items())
with self._lock:
child = self._children.get(key)
if child is None:
child = PythonWorkerPool(
self.python,
size=self.size,
events=self.events,
env={**self.env, **extra},
)
child._generation = self._generation
child.start()
self._children[key] = child
return child
def stop(self) -> None:
self._stopped = True
self._generation += 1
for child in list(self._children.values()):
child.stop()
for worker in list(self._running.values()):
worker.kill()
for slot in self._drain():
@@ -281,6 +317,9 @@ class PythonWorkerPool:
if slot is not None:
slot.kill()
self._idle.put(None)
children = list(self._children.values())
for child in children:
child.respawn_all()
def _drain(self) -> list[_Worker | None]:
slots = []
@@ -575,11 +614,14 @@ class PythonWorkerPool:
"""Stop a node that is running now. False when there was nothing to stop."""
with self._lock:
worker = self._running.get((run_id, node_id))
if worker is None:
return False
worker.cancelled = True
worker.kill()
return True
children = list(self._children.values())
if worker is not None:
worker.cancelled = True
worker.kill()
return True
# A node with declared resources runs in a pool of its own, and it is
# the same node to whoever pressed cancel.
return any(child.cancel(node_id, run_id) for child in children)
def cancel_run(self, run_id: str) -> int:
"""Stop every node this run has in a worker right now."""
@@ -592,7 +634,8 @@ class PythonWorkerPool:
for worker in workers:
worker.cancelled = True
worker.kill()
return len(workers)
children = list(self._children.values())
return len(workers) + sum(child.cancel_run(run_id) for child in children)
def _publish(self, event: dict[str, Any]) -> None:
if self.events is not None: