"""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}