"""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: the books for one machine, which hand out what is free and take it back. :func:`derive_env` is the second: what a node got is handed to the library as the environment it reads at import, because that is the only moment those knobs are read. One accountant is one machine. Which machine a node goes to, and the waiting when none of them has room, is :mod:`fluksio.flow.placement` — the books here never block, so a caller holding several of them can ask each in turn. 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. """ from __future__ import annotations import logging import os import re import threading from collections.abc import Callable from dataclasses import dataclass from sqlmodel import Session from fluksio.flow.schemas import Resources logger = logging.getLogger(__name__) #: What a flavor may be called. Dashes allowed, unlike a flow or node name — #: "gpu-small" reads better than "gpu_small" on a dropdown. FLAVOR_NAME = re.compile(r"^[a-z][a-z0-9_-]*$") #: 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, ...] = () ram_mb: int = 0 class ResourceAccountant: """The books for one machine: what it holds, and what is free of it. Nothing here blocks. A node that has to wait waits in the placer, which holds one of these per machine and tries each in turn — a lock per machine could not answer "is there room anywhere", and a node waiting inside one machine's books could not be woken by another machine attaching. ``on_release`` is how it tells the placer something came free. It is called outside this object's lock, deliberately: the placer takes its own lock there, and the two are only ever taken in that one order. """ def __init__( self, cpus: int = 0, gpus: int = 0, ram_mb: int | None = None, on_release: Callable[[], None] | None = None, ) -> None: self.cpus = cpus or machine_cpus() self.gpus = max(0, gpus) #: None is memory this machine did not report, which is allowed: it is #: then not accounted rather than accounted as zero, since zero would #: refuse every node that asks for any. self.ram_mb = ram_mb self.on_release = on_release self._free_cpus = self.cpus self._free_gpus = list(range(self.gpus)) self._free_ram = ram_mb or 0 self._lock = threading.Lock() def fits(self, cpus: int, gpus: int, ram_mb: int = 0) -> bool: """Whether this machine could ever grant that much, busy or not. The question a placer asks before it decides to wait: waiting for a card that is merely busy is queueing, and waiting for one the machine does not have is a hung run. """ if cpus > self.cpus or gpus > self.gpus: return False return self.ram_mb is None or ram_mb <= self.ram_mb def try_take(self, cpus: int, gpus: int, ram_mb: int = 0) -> Allocation | None: """Take this much if it is free right now, or answer that it is not.""" with self._lock: if self._free_cpus < cpus or len(self._free_gpus) < gpus: return None if self.ram_mb is not None and self._free_ram < ram_mb: return None self._free_cpus -= cpus taken = tuple(self._free_gpus.pop(0) for _ in range(gpus)) if self.ram_mb is not None: self._free_ram -= ram_mb return Allocation(cpus=cpus, gpus=taken, ram_mb=ram_mb) def give_back(self, allocation: Allocation) -> None: """Return what an execution held, however it ended.""" with self._lock: self._free_cpus += allocation.cpus self._free_gpus.extend(allocation.gpus) self._free_gpus.sort() if self.ram_mb is not None: self._free_ram += allocation.ram_mb if self.on_release is not None: self.on_release() def snapshot(self) -> dict[str, object]: """What this machine holds and what is free of it.""" with self._lock: free_cpus, free_gpus, free_ram = ( self._free_cpus, len(self._free_gpus), self._free_ram, ) return { "cpus": {"total": self.cpus, "free": free_cpus}, "gpus": {"total": self.gpus, "free": free_gpus}, "ram_mb": ( None if self.ram_mb is None else {"total": self.ram_mb, "free": free_ram} ), } # ----------------------------------------------------------------------------- # 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 class UnknownFlavor(ValueError): """A node asks for a size that is not stored here.""" def resolve_flavor(wanted: Resources) -> Resources: """The concrete numbers behind a declaration. Read when the node is built rather than stored on it, so editing a flavor changes what the next run gets. A node that names one that has been deleted is an error rather than a default: running a training step against a size nobody chose is worse than a node that says what is wrong with it. """ if not wanted.flavor: return wanted # Imported here: this module is the books, and the books have no business # knowing about the database until somebody asks for a stored size. from fluksio.core.db import engine from fluksio.models import Flavor with Session(engine) as session: row = session.get(Flavor, wanted.flavor) if row is None: raise UnknownFlavor( f"flavor '{wanted.flavor}' does not exist — `fluksio flavors` lists them" ) # Built fresh rather than copied, so the flavor-and-numbers check runs on # the result and this cannot quietly produce something invalid. return Resources( cpus=row.cpus, gpus=row.gpus, ram=row.ram, env=wanted.env, duration_s=wanted.duration_s, ) 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}