"""Admission control from live memory and CPU limits, not a fixed worker count.""" import math import os from dataclasses import dataclass from pathlib import Path import psutil RESERVE = 0.10 @dataclass(frozen=True) class Resources: ram_free: int vram_free: int cpu_capacity: float def _read_limit(path: str) -> int | None: try: text = Path(path).read_text().strip() return None if text == "max" else int(text) except (OSError, ValueError): return None def snapshot(vram_free: int) -> Resources: ram = psutil.virtual_memory().available limit = _read_limit("/sys/fs/cgroup/memory.max") used = _read_limit("/sys/fs/cgroup/memory.current") if limit is not None and used is not None: ram = min(ram, max(0, limit - used)) cpus = float(len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1) try: quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split() if quota != "max": cpus = min(cpus, int(quota) / int(period)) except (OSError, ValueError): pass return Resources(ram, vram_free, cpus) def parallelism(resources: Resources, ram_per_task: int, vram_per_task: int, pending: int) -> int: if pending < 1 or min(ram_per_task, vram_per_task) <= 0: raise ValueError("Pending work and measured memory costs must be positive.") ram_slots = int(resources.ram_free * (1 - RESERVE)) // ram_per_task gpu_slots = int(resources.vram_free * (1 - RESERVE)) // vram_per_task if min(ram_slots, gpu_slots) < 1: raise MemoryError("Even one texture task cannot fit with the 10% memory reserve; reduce tile size.") cpu_slots = max(1, math.floor(resources.cpu_capacity * (1 - RESERVE))) return max(1, min(pending, ram_slots, gpu_slots, cpu_slots))