stroblme
e37aa2bdd7
Per the Microcopy rule: size and duration errors, the flavor errors, run parameter errors, the Resources field descriptions, 401/409/400 details, secret errors, the websocket node description and its on_connect setting. The generated client follows.
327 lines
12 KiB
Python
327 lines
12 KiB
Python
"""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 glob
|
|
import logging
|
|
import os
|
|
import re
|
|
import threading
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
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
|
|
|
|
#: The same for memory, in megabytes.
|
|
ENGINE_RAM_RESERVE_MB = 1024
|
|
|
|
#: Where a container finds its own memory limit (cgroup v2).
|
|
CGROUP_MEMORY_MAX = Path("/sys/fs/cgroup/memory.max")
|
|
|
|
#: How many times the core count a concurrency limit may reach before startup
|
|
#: warns. Above one on purpose: a node waiting on the network holds no core.
|
|
OVERSUBSCRIBE_FACTOR = 4
|
|
|
|
|
|
def machine_cpus() -> int:
|
|
"""Cores a node may be given here, when nobody said."""
|
|
return max(1, (os.cpu_count() or 1) - ENGINE_RESERVE)
|
|
|
|
|
|
def warn_oversubscribed(limits: dict[str, int]) -> None:
|
|
"""Warn, never clamp, when a concurrency limit is far past the cores.
|
|
|
|
A high number can be right for nodes that wait on devices, but one far
|
|
past the machine is more often a zero too many.
|
|
"""
|
|
cores = os.cpu_count() or 1
|
|
high = [
|
|
f"{name}={value}"
|
|
for name, value in limits.items()
|
|
if value > OVERSUBSCRIBE_FACTOR * cores
|
|
]
|
|
if high:
|
|
logger.warning(
|
|
"Concurrency above %dx this machine's %d cores: %s",
|
|
OVERSUBSCRIBE_FACTOR,
|
|
cores,
|
|
", ".join(high),
|
|
)
|
|
|
|
|
|
def visible_gpus() -> list[str] | None:
|
|
"""The cards the engine's own ``CUDA_VISIBLE_DEVICES`` lists, if it is set.
|
|
|
|
Indices, UUIDs or MIG instances, kept as written: a child is handed the
|
|
same names, so whatever the driver took here it takes there.
|
|
"""
|
|
mask = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
if mask is None:
|
|
return None
|
|
return [part.strip() for part in mask.split(",") if part.strip()]
|
|
|
|
|
|
def machine_gpus() -> int:
|
|
"""Cards a node may be given here, when nobody said.
|
|
|
|
The engine's own ``CUDA_VISIBLE_DEVICES`` where it has one, since a card
|
|
the process can see but was not given is not its to hand out. Otherwise
|
|
the device nodes the NVIDIA driver creates, counted — no vendor tool, so
|
|
the one dependency does not become two. It answers for the case that
|
|
actually bit people (a box with cards, an engine told nothing, every GPU
|
|
node clamped to zero and running at once) and for nothing else: another
|
|
vendor, or a card this process cannot see, is still `--gpus N`.
|
|
"""
|
|
visible = visible_gpus()
|
|
if visible is not None:
|
|
return len(visible)
|
|
return len(glob.glob("/dev/nvidia[0-9]*"))
|
|
|
|
|
|
def machine_ram_mb() -> int | None:
|
|
"""Megabytes a node may be given here, when nobody said.
|
|
|
|
The box's memory, or the container's limit where that is lower — the
|
|
host's total is memory the engine would be killed for using. None where
|
|
neither can be asked, which leaves memory unaccounted rather than zero.
|
|
"""
|
|
try:
|
|
total = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") // 2**20
|
|
except (AttributeError, ValueError, OSError):
|
|
return None
|
|
try:
|
|
# "max" when the container has no limit of its own.
|
|
limit = CGROUP_MEMORY_MAX.read_text().strip()
|
|
except OSError:
|
|
limit = ""
|
|
if limit.isdigit():
|
|
total = min(total, int(limit) // 2**20)
|
|
return max(1, total - ENGINE_RAM_RESERVE_MB)
|
|
|
|
|
|
@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, local: bool = True
|
|
) -> 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.
|
|
|
|
``local`` is whether the allocation is this machine's. A child's
|
|
``CUDA_VISIBLE_DEVICES`` replaces the engine's rather than narrowing it, so
|
|
here a card is named the way the engine's own list names it: under ``=2,3``
|
|
card 0 is device 2. A remote worker does the same against its own list.
|
|
"""
|
|
env = {var: str(allocation.cpus) for var in THREAD_VARS}
|
|
if wanted.gpus:
|
|
names = (visible_gpus() if local else None) or []
|
|
env["CUDA_VISIBLE_DEVICES"] = ",".join(
|
|
names[index] if index < len(names) else 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}
|