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