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