studio/resources.py

"""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))