app/backend/fluksio/flow/resources.py

"""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 threading
from collections.abc import Callable
from dataclasses import dataclass

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, ...] = ()
    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


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}