"""Keep the first run from being killed by its own compiler.

The text weights of this checkpoint are NVFP4.  On a Blackwell card vLLM
reaches FlashInfer's CUTLASS FP4 GEMM, and on a machine that has never run it
before, that kernel is compiled on the spot - eighteen CUTLASS translation
units, launched by ninja at whatever parallelism the host reports.

Each of those units costs several gigabytes of host memory while nvcc expands
the templates.  Measured on a rented RTX 5090 with 96 visible cores and an 85
GiB memory limit: the default parallelism was killed by the out-of-memory
killer (ninja reported ``FAILED: [code=137]``), and the engine never started.
Capped at four jobs the same build finished in about four minutes and cached
its result, after which start-up is immediate.

A person with a 5090 and 32 GiB of system memory would hit exactly that wall on
their first command, and the failure gives no hint about its cause.  So the
model sets a build parallelism it can afford - and only when the user has not
chosen one.
"""

from __future__ import annotations

import os
from pathlib import Path

REPAIR_ID = "ZENIT_BUILD_PARALLELISM_GUARD_V1"

# Host memory one CUTLASS FP4 translation unit needs while nvcc runs.  Derived
# from the measurement above: eighteen concurrent jobs exceeded 85 GiB, four
# fitted comfortably.  Kept deliberately generous - being slow to compile once
# is cheap, being killed is not.
BYTES_PER_JOB = 6 * 1024**3

# Beyond this, extra jobs buy little: the build is eighteen units and the tail
# is dominated by the largest few.
MAX_PARALLEL_JOBS = 8


def _read_int(path: str) -> int | None:
    try:
        text = Path(path).read_text(encoding="utf-8").strip()
    except OSError:
        return None
    if text in ("max", ""):
        return None
    try:
        return int(text)
    except ValueError:
        return None


def container_memory_limit() -> int | None:
    """Memory this process may use, honouring cgroup limits.

    ``/proc/meminfo`` reports the whole machine, which is the wrong number
    inside a container: the compiler is killed at the cgroup limit, not at the
    host's.
    """
    for limit_path, usage_path in (
        ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory.current"),
        (
            "/sys/fs/cgroup/memory/memory.limit_in_bytes",
            "/sys/fs/cgroup/memory/memory.usage_in_bytes",
        ),
    ):
        limit = _read_int(limit_path)
        if limit is None or limit > (1 << 60):  # "no limit" sentinel
            continue
        usage = _read_int(usage_path) or 0
        return max(limit - usage, 0)
    return None


def host_available_memory() -> int | None:
    try:
        for line in Path("/proc/meminfo").read_text(encoding="utf-8").splitlines():
            if line.startswith("MemAvailable:"):
                return int(line.split()[1]) * 1024
    except (OSError, IndexError, ValueError):
        return None
    return None


def available_memory() -> int | None:
    """The smaller of the cgroup allowance and what the host reports free."""
    candidates = [
        value
        for value in (container_memory_limit(), host_available_memory())
        if value is not None
    ]
    return min(candidates) if candidates else None


def choose_job_count(
    available_bytes: int | None, cpu_count: int | None = None
) -> int | None:
    """How many compiler jobs this machine can afford; ``None`` if unknown."""
    if available_bytes is None:
        return None
    cpus = cpu_count if cpu_count is not None else (os.cpu_count() or 1)
    affordable = int(available_bytes // BYTES_PER_JOB)
    return max(1, min(cpus, affordable, MAX_PARALLEL_JOBS))


def install_build_parallelism_guard() -> dict[str, str] | None:
    """Set ``MAX_JOBS`` and ``NVCC_THREADS`` when the caller has not.

    Returns the values it set, or ``None`` when it left the environment alone.
    """
    if os.environ.get("MAX_JOBS"):
        return None  # the user made a choice; it stands
    jobs = choose_job_count(available_memory())
    if jobs is None:
        return None
    applied = {"MAX_JOBS": str(jobs)}
    os.environ["MAX_JOBS"] = str(jobs)
    if not os.environ.get("NVCC_THREADS"):
        # Two device-code threads per job keeps the memory estimate honest.
        os.environ["NVCC_THREADS"] = "2"
        applied["NVCC_THREADS"] = "2"
    return applied


__all__ = [
    "BYTES_PER_JOB",
    "MAX_PARALLEL_JOBS",
    "REPAIR_ID",
    "available_memory",
    "choose_job_count",
    "container_memory_limit",
    "host_available_memory",
    "install_build_parallelism_guard",
]