diff --git a/backend/fluksio/api/routes/workers.py b/backend/fluksio/api/routes/workers.py index 9ed85cb..88c4b02 100644 --- a/backend/fluksio/api/routes/workers.py +++ b/backend/fluksio/api/routes/workers.py @@ -79,6 +79,20 @@ def read_workers(request: Request) -> Any: ] +@router.get("/resources", dependencies=[Depends(get_current_user)]) +def read_resources(request: Request) -> Any: + """What this machine has free, and which nodes are queued for it. + + A node waiting its turn looks exactly like a node that has hung — the run + sits at `running` and says nothing — so what is waiting, and for what, has + to be readable somewhere. + """ + accountant = getattr(request.app.state, "resources", None) + if accountant is None: + raise HTTPException(status_code=503, detail="Resources are not accounted here") + return accountant.snapshot() + + @router.post( "/tokens", response_model=TokenIssued, diff --git a/backend/fluksio/core/config.py b/backend/fluksio/core/config.py index b1ca365..39a8012 100644 --- a/backend/fluksio/core/config.py +++ b/backend/fluksio/core/config.py @@ -102,6 +102,15 @@ class Settings(BaseSettings): # fails fast, and a slow one is left to finish. Set it where silence means # stuck rather than working. FLOW_NODE_TIMEOUT: float = 0.0 + # Cores nodes may be given, for the ones that declare `resources`. 0 works + # it out: every core but two, which are what keeps the engine's own event + # loop answering while the machine is busy. Nodes that declare nothing are + # not accounted against it — they only get its fair share as a thread cap. + FLOW_CPUS: int = 0 + # GPUs on this machine. Not detected, because detecting it means depending + # on the vendor's tooling: say how many there are and each is held by one + # node at a time. + FLOW_GPUS: int = 0 # How long the engine's own metrics, events and run records are kept. OBS_RETENTION_DAYS: int = 30 # Without a Redis host the engine keeps its state in memory. diff --git a/backend/fluksio/flow/controller.py b/backend/fluksio/flow/controller.py index f321936..d4daaf0 100644 --- a/backend/fluksio/flow/controller.py +++ b/backend/fluksio/flow/controller.py @@ -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 diff --git a/backend/fluksio/flow/resources.py b/backend/fluksio/flow/resources.py new file mode 100644 index 0000000..413906a --- /dev/null +++ b/backend/fluksio/flow/resources.py @@ -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} diff --git a/backend/fluksio/flow/schemas.py b/backend/fluksio/flow/schemas.py index eace210..535e3b9 100644 --- a/backend/fluksio/flow/schemas.py +++ b/backend/fluksio/flow/schemas.py @@ -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 diff --git a/backend/fluksio/flow/workers.py b/backend/fluksio/flow/workers.py index 985b2a7..06180bf 100644 --- a/backend/fluksio/flow/workers.py +++ b/backend/fluksio/flow/workers.py @@ -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: diff --git a/backend/fluksio/main.py b/backend/fluksio/main.py index 587ceaa..7456ddd 100644 --- a/backend/fluksio/main.py +++ b/backend/fluksio/main.py @@ -31,6 +31,7 @@ from fluksio.flow.pipeline import ValueSource from fluksio.flow.plugins import load_plugins from fluksio.flow.queue import MemoryWorkQueue, RedisWorkQueue, WorkQueue from fluksio.flow.remote import RemoteWorkerHub +from fluksio.flow.resources import ResourceAccountant, fair_share_env from fluksio.flow.runs import RUN_STATE_TTL, RunService from fluksio.flow.secrets import init_secrets from fluksio.flow.state import MemoryState, RedisState, StateBackend @@ -116,13 +117,24 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # not something anyone wrote, so it has no business in the git repository. artifacts = ArtifactStore(settings.FLOWS_DIR.parent / "artifacts") app.state.artifact_store = artifacts + accountant = ResourceAccountant( + cpus=settings.FLOW_CPUS, gpus=settings.FLOW_GPUS, events=event_bus + ) + app.state.resources = accountant pool = PythonWorkerPool( python=modules.venv_python(), size=settings.FLOW_MAX_WORKERS, events=event_bus, # A worker in this container writes to the store directly; a remote one # is given a URL instead. Node code calls the same two functions. - env={ARTIFACT_DIR_ENV: str(artifacts.root)}, + env={ + ARTIFACT_DIR_ENV: str(artifacts.root), + # Every slot can be busy at once, so a worker left to size its own + # thread pool to the machine means as many processes as there are + # slots, each believing it has the whole of it. A node that says + # what it needs overrides this; one that says nothing gets a share. + **fair_share_env(accountant.cpus, settings.FLOW_MAX_WORKERS), + }, ) pool.start() app.state.worker_pool = pool @@ -138,6 +150,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: alerts=alerts, workers=pool, remote=worker_hub, + resources=accountant, ) app.state.flow_controller = controller # A "dashboard" alert channel puts its alert into the graph. Bound here diff --git a/backend/fluksio/sdk/__init__.py b/backend/fluksio/sdk/__init__.py index 5929cca..e41f1aa 100644 --- a/backend/fluksio/sdk/__init__.py +++ b/backend/fluksio/sdk/__init__.py @@ -178,6 +178,7 @@ class NodeSpec: device: str | None, device_policy: str, cache: bool = True, + resources: dict[str, Any] | None = None, ) -> None: self.fn = fn self.id = id @@ -190,6 +191,7 @@ class NodeSpec: self.device = device self.device_policy = device_policy self.cache = cache + self.resources = resources def rebind( self, *, id: str = "", wire: dict[str, str] | None = None, **settings: Any @@ -221,6 +223,7 @@ class NodeSpec: device=self.device, device_policy=self.device_policy, cache=self.cache, + resources=self.resources, ) @@ -252,6 +255,7 @@ def node( device: str | None = None, device_policy: str = "require", cache: bool = True, + resources: dict[str, Any] | None = None, ) -> Callable[[F], F]: """Mark a function as a node, declaring its ports. @@ -271,9 +275,29 @@ def node( ``timeout`` is seconds of silence — a yield or an emit resets it — after which the node is stopped. Left out it inherits the engine's default, which is no limit; set one where silence means stuck rather than working. + + ``resources`` says how much of the machine one execution takes, for a node + heavy enough that the answer is not "a share of it":: + + @node(..., resources={"cpus": 4}) + @node(..., resources={"gpus": 1, + "env": {"XLA_PYTHON_CLIENT_MEM_FRACTION": "0.9"}}) + + The engine holds that much while the node runs, so the rest of them wait + their turn rather than oversubscribing the box, and the worker is started + with thread limits matching what it was given — which is the only way a + library that sizes itself to every visible core is told otherwise. ``env`` + is for the tuning the engine must not invent, such as ``XLA_FLAGS``. + Declaring nothing is the default and changes nothing. """ if device_policy not in ("require", "prefer"): raise SyncError("device_policy is 'require' or 'prefer'") + if resources is not None: + unknown = sorted(set(resources) - {"cpus", "gpus", "env"}) + if unknown: + raise SyncError( + f"resources={{'{unknown[0]}': ...}} is not one of cpus, gpus, env" + ) if timeout is not None and timeout < 0: raise SyncError( f"node timeout must be 0 or more (0 disables it), got {timeout}" @@ -294,6 +318,7 @@ def node( device=device, device_policy=device_policy, cache=cache, + resources=dict(resources) if resources is not None else None, ) _check_signature(spec) fn.__fluksio__ = spec # type: ignore[attr-defined] @@ -499,6 +524,7 @@ def _check(spec: NodeSpec, mode: str) -> dict[str, Any]: "device": spec.device, "device_policy": spec.device_policy, "cache": spec.cache, + "resources": spec.resources, } diff --git a/backend/tests/flow/test_resources.py b/backend/tests/flow/test_resources.py new file mode 100644 index 0000000..30551ad --- /dev/null +++ b/backend/tests/flow/test_resources.py @@ -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 diff --git a/backend/tests/flow/test_workers.py b/backend/tests/flow/test_workers.py index aa5dc1a..a838d2c 100644 --- a/backend/tests/flow/test_workers.py +++ b/backend/tests/flow/test_workers.py @@ -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 diff --git a/docs/getting-started/data-science.md b/docs/getting-started/data-science.md index 876ca13..6684ec2 100644 --- a/docs/getting-started/data-science.md +++ b/docs/getting-started/data-science.md @@ -262,6 +262,54 @@ stops your own code from running. )]) ``` +### Declaring what a node needs + +A training node is not like the rest of a flow. NumPy, JAX and PyTorch each +size their thread pool to every core they can see, so a few of them at once +oversubscribe the machine badly enough to starve the engine itself — the API +stops answering and every client waiting on a run dies with it. On a GPU the +same shape deadlocks: two processes each preallocating most of the card sit at +zero utilisation with nothing failing and nothing to read. + +Say what one execution takes, and the engine holds it: + +```python +@node(..., resources={"cpus": 2, "env": {"XLA_FLAGS": "--xla_cpu_multi_thread_eigen=false"}}) +def fit(dataset, lr, epochs=25): + ... + +@node(..., resources={"gpus": 1, "env": {"XLA_PYTHON_CLIENT_MEM_FRACTION": "0.9"}}) +def finetune(checkpoint): + ... +``` + +Two things follow. The node waits its turn rather than starting alongside more +of itself than the machine has room for — the same waiting it already does for +a worker. And the worker it runs in is *started* with thread limits matching +what it was given, because that is the only moment a numerical library reads +them; a GPU node is told which card is its through `CUDA_VISIBLE_DEVICES`, and +nothing else is given that card while it runs. + +`env` is for the tuning the engine must not invent. `XLA_FLAGS` is one composed +string whose contents depend on the version you have installed, so writing it +for you would silently replace whatever you had put there. + +Declaring nothing is the default and is right for most nodes — a poll, a +threshold, a message on its way somewhere. Those share the engine's worker pool +and are given a fair share of `FLOW_CPUS` as a thread cap, which is what stops +several of them at once from each sizing to the whole box. + +Ask what is free, and what is waiting for it, at +`GET /api/v1/workers/resources` — a node queued for cores looks exactly like a +node that has hung unless you can see the queue. + +!!! note "Accounted, not enforced" + + Nothing stops a node that ignores its declaration; the numbers are + bookkeeping plus the environment its libraries read. Scheduling across + several machines, named hardware flavours and real enforcement are the + next steps, not this one. + ## Say which nodes make a flow Membership is a list, not a directory layout: the functions can live wherever diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 67a0b15..e00d15f 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -125,8 +125,18 @@ warning into a refusal to start. | `FLOW_MAX_WORKERS` | `4` | node-code subprocesses run in parallel | | `FLOW_MAX_CASCADES` | `4` | cascades in flight at once; throughput is this over the mean cascade time, so raise it where nodes wait on a network rather than a CPU | | `FLOW_NODE_TIMEOUT` | `0` | seconds a node may be silent, unless it sets its own; 0 is no limit | +| `FLOW_CPUS` | `0` | cores nodes that declare `resources` may be given; 0 works it out as every core but two, which are what keeps the engine answering while the machine is busy | +| `FLOW_GPUS` | `0` | GPUs on this machine, each held by one node at a time. Not detected — say how many there are | | `OBS_RETENTION_DAYS` | `30` | how long metrics, events and run records are kept | +A node that declares nothing is not accounted against `FLOW_CPUS`; it runs on +the shared pool and is given `FLOW_CPUS / FLOW_MAX_WORKERS` as a thread cap, so +several at once cannot each size themselves to the whole machine. Setting +`OMP_NUM_THREADS` (or any of its siblings) on the engine yourself overrides +that default. What is free, and which nodes are queued for it, is +`GET /api/v1/workers/resources`. See +[declaring resources](../getting-started/data-science.md#declaring-what-a-node-needs). + ## Agents | Variable | Default | Notes | diff --git a/uv.lock b/uv.lock index 698564d..e3ae1a9 100644 --- a/uv.lock +++ b/uv.lock @@ -869,7 +869,7 @@ wheels = [ [[package]] name = "fluksio" -version = "0.1.3" +version = "0.1.4" source = { editable = "backend" } dependencies = [ { name = "aiomqtt" }, @@ -948,7 +948,7 @@ dev = [ [[package]] name = "fluksio-worker" -version = "0.1.3" +version = "0.1.4" source = { editable = "worker" } dependencies = [ { name = "websockets" },