runtime/src/resynthesis/release_cli.py

"""Minimal public command line for prompting Nucleus Resynthesis Release 188."""

from __future__ import annotations

import argparse
import fcntl
import os
import subprocess
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from pathlib import Path
from typing import Final, Protocol, TextIO

import torch

_LAB_GPU_LOCK_ROOT: Final[Path] = Path("/tmp")
_BENCHMARK_GPU_LOCK_ROOT: Final[Path] = Path(
    "/tmp/nnf_benchmark_gpu_locks"
)
_MINIMUM_GPU_HEADROOM_BYTES: Final[int] = 48 * 1024**3


class _ReleaseSolveModel(Protocol):
    def solve(self, prompt: str) -> str:
        """Return the model-owned decoded response for one prompt."""


def _nvidia_identity_rows_boundary() -> tuple[tuple[str, int], ...]:
    """Read physical GPU UUIDs at the external host-discovery boundary."""

    try:
        completed = subprocess.run(
            (
                "nvidia-smi",
                "--query-gpu=index,uuid",
                "--format=csv,noheader,nounits",
            ),
            check=True,
            capture_output=True,
            text=True,
        )
    except (FileNotFoundError, subprocess.SubprocessError):
        return ()
    rows: list[tuple[str, int]] = []
    for line in completed.stdout.splitlines():
        fields = tuple(part.strip() for part in line.split(",", maxsplit=1))
        if len(fields) != 2 or not fields[0].isdigit() or not fields[1]:
            continue
        rows.append((fields[1], int(fields[0])))
    return tuple(rows)


def _visible_physical_gpu_indices_boundary(
    device_count: int,
) -> tuple[int, ...]:
    """Map process-local CUDA ordinals to physical lock-file identities."""

    visible_text = os.environ.get("CUDA_VISIBLE_DEVICES")
    if visible_text is None:
        return tuple(range(device_count))
    visible_tokens = tuple(
        token.strip() for token in visible_text.split(",") if token.strip()
    )
    if (
        not visible_tokens
        or visible_tokens == ("-1",)
        or len(visible_tokens) != device_count
    ):
        return ()
    identity_rows = (
        ()
        if all(token.isdigit() for token in visible_tokens)
        else _nvidia_identity_rows_boundary()
    )
    physical_indices: list[int] = []
    for token in visible_tokens:
        if token.isdigit():
            physical_indices.append(int(token))
            continue
        matches = tuple(
            physical_index
            for uuid, physical_index in identity_rows
            if uuid == token or uuid.startswith(token)
        )
        if len(matches) != 1:
            return ()
        physical_indices.append(matches[0])
    if len(set(physical_indices)) != len(physical_indices):
        return ()
    return tuple(physical_indices)


def _try_acquire_gpu_locks_boundary(
    physical_index: int,
) -> tuple[TextIO, TextIO] | None:
    """Acquire both standard exclusive GPU locks without waiting."""

    _LAB_GPU_LOCK_ROOT.mkdir(parents=True, exist_ok=True)
    _BENCHMARK_GPU_LOCK_ROOT.mkdir(parents=True, exist_ok=True)
    paths = (
        _LAB_GPU_LOCK_ROOT / f"nnf_gpu_{physical_index}.lock",
        _BENCHMARK_GPU_LOCK_ROOT / f"gpu_{physical_index}.lock",
    )
    handles: list[TextIO] = []
    try:
        for path in paths:
            handle = path.open("a+", encoding="utf-8")
            try:
                fcntl.flock(
                    handle.fileno(),
                    fcntl.LOCK_EX | fcntl.LOCK_NB,
                )
            except BlockingIOError:
                handle.close()
                return None
            handles.append(handle)
    finally:
        if len(handles) != len(paths):
            for cleanup_handle in reversed(handles):
                fcntl.flock(cleanup_handle.fileno(), fcntl.LOCK_UN)
                cleanup_handle.close()
    return handles[0], handles[1]


def _release_gpu_locks_boundary(handles: tuple[TextIO, TextIO]) -> None:
    for handle in reversed(handles):
        fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
        handle.close()


@contextmanager
def _release_compute_lane_boundary(
    *,
    minimum_free_bytes: int = _MINIMUM_GPU_HEADROOM_BYTES,
) -> Iterator[torch.device]:
    """Hold one clean GPU lane for load and solve, or fall back to CPU."""

    if isinstance(minimum_free_bytes, bool) or minimum_free_bytes < 1:
        raise ValueError("Release 188 minimum GPU headroom must be positive")
    if not torch.cuda.is_available():
        yield torch.device("cpu")
        return
    device_count = torch.cuda.device_count()
    physical_indices = _visible_physical_gpu_indices_boundary(device_count)
    if len(physical_indices) != device_count:
        yield torch.device("cpu")
        return
    candidates = tuple(
        sorted(
            (
                (
                    torch.cuda.mem_get_info(local_index)[0],
                    physical_index,
                    local_index,
                )
                for local_index, physical_index in enumerate(physical_indices)
            ),
            reverse=True,
        )
    )
    for observed_free_bytes, physical_index, local_index in candidates:
        if observed_free_bytes < minimum_free_bytes:
            continue
        handles = _try_acquire_gpu_locks_boundary(physical_index)
        if handles is None:
            continue
        try:
            free_bytes_after_lock = torch.cuda.mem_get_info(local_index)[0]
            if free_bytes_after_lock < minimum_free_bytes:
                continue
            torch.cuda.set_device(local_index)
            yield torch.device("cuda", local_index)
            return
        finally:
            _release_gpu_locks_boundary(handles)
    yield torch.device("cpu")


def _load_release_model_boundary(
    release_path: Path,
    *,
    device: torch.device,
) -> _ReleaseSolveModel:
    from resynthesis.release_model import load_release_188_model

    return load_release_188_model(release_path, device=device)


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="resynthesis",
        description="Prompt Nucleus Resynthesis Release 188.",
    )
    commands = parser.add_subparsers(dest="command", required=True)
    solve = commands.add_parser(
        "solve",
        help="Generate one model-owned response.",
    )
    solve.add_argument(
        "--release",
        required=True,
        type=Path,
        help="Path to runtime/model.json.",
    )
    solve.add_argument("prompt", help="Prompt text.")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    """Run the minimal inference-only Release 188 command surface."""

    arguments = _parser().parse_args(argv)
    if arguments.command != "solve":
        raise RuntimeError("Release 188 command dispatch differs")
    release_path = arguments.release
    prompt = arguments.prompt
    if not isinstance(release_path, Path) or not isinstance(prompt, str):
        raise RuntimeError("Release 188 solve arguments are malformed")
    with _release_compute_lane_boundary() as device:
        model = _load_release_model_boundary(
            release_path,
            device=device,
        )
        response = model.solve(prompt)
    print(response)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())