OBLITERATUS/obliteratus/gpu_lifecycle.py

"""Local, credential-free GPU lifecycle publication for host supervisors."""

from __future__ import annotations

import atexit
from dataclasses import dataclass
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import threading
import uuid
from typing import Callable


SCHEMA_VERSION = 1


class AdmissionError(RuntimeError):
    """The local supervisor did not grant correlated GPU admission."""


@dataclass(frozen=True)
class MemoryUsage:
    """Process-visible accelerator memory in bytes."""

    allocated_bytes: int = 0
    reserved_bytes: int = 0
    device_count: int = 0


class GpuLifecyclePublisher:
    """Publish ordered lifecycle events into an operator-owned runtime directory."""

    def __init__(
        self,
        runtime_dir: str | Path | None,
        *,
        heartbeat_seconds: float = 15.0,
        admission_timeout_seconds: float = 30.0,
        admission_poll_seconds: float = 0.05,
        clock: Callable[[], datetime] | None = None,
        run_id: str | None = None,
    ) -> None:
        configured_dir = Path(runtime_dir) if runtime_dir else None
        if configured_dir is not None and configured_dir.is_symlink():
            raise ValueError("GPU lifecycle runtime directory must not be a symlink")
        self._dir = configured_dir.resolve() if configured_dir else None
        self._heartbeat_seconds = max(1.0, float(heartbeat_seconds))
        self._admission_timeout_seconds = max(0.01, float(admission_timeout_seconds))
        self._admission_poll_seconds = max(0.001, float(admission_poll_seconds))
        self._clock = clock or (lambda: datetime.now(timezone.utc))
        self._run_id = run_id or str(uuid.uuid4())
        self._lock = threading.RLock()
        self._sequence = 0
        self._model_id: str | None = None
        self._phase = "released"
        self._last_memory = MemoryUsage()
        self._lease_id: str | None = None
        self._granted_vram_bytes = 0
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        if self._dir is not None:
            if not self._dir.is_dir():
                raise ValueError("GPU lifecycle runtime directory must already exist")
        atexit.register(self.release, reason="process_exit")

    @property
    def enabled(self) -> bool:
        return self._dir is not None

    def loading(self, model_id: str) -> dict | None:
        with self._lock:
            self._model_id = str(model_id)
            self._phase = "intent_published"
            request = self._publish("loading")
            if request is None:
                return None
            self._wait_for_admission(request)
            self._phase = "admission_granted"
            self._publish("admission_granted")
            self._phase = "allocation_started"
            return self._publish("allocation_started")

    def resize(self, memory: MemoryUsage) -> dict | None:
        with self._lock:
            if (
                self.enabled
                and memory.reserved_bytes > self._granted_vram_bytes
            ):
                self._phase = "admission_lost"
                self._publish("admission_lost", reason="reservation_exceeded")
                raise AdmissionError("measured VRAM exceeds the supervisor grant")
            self._last_memory = memory
            return self._publish("resize")

    def ready(self, memory: MemoryUsage | None = None) -> dict | None:
        with self._lock:
            if memory is not None:
                self._last_memory = memory
            self._phase = "ready"
            event = self._publish("ready")
            self._start_heartbeat()
            return event

    def heartbeat(self) -> dict | None:
        with self._lock:
            if self._phase != "ready":
                return None
            return self._publish("heartbeat")

    def release(self, *, reason: str = "unload") -> dict | None:
        with self._lock:
            if self._phase == "released":
                return None
            self._phase = "released"
            self._stop.set()
            event = self._publish("release", reason=reason)
            self._model_id = None
            self._last_memory = MemoryUsage()
            self._lease_id = None
            self._granted_vram_bytes = 0
            return event

    def _wait_for_admission(self, request: dict) -> None:
        acknowledgement = self._dir / "ack.json"  # type: ignore[operator]
        deadline = threading.Event()
        remaining = self._admission_timeout_seconds
        while remaining > 0:
            started = datetime.now(timezone.utc)
            try:
                payload = json.loads(acknowledgement.read_text(encoding="utf-8"))
            except (FileNotFoundError, OSError, json.JSONDecodeError):
                payload = None
            if self._valid_ack(payload, request):
                if payload["decision"] == "deny":
                    self._phase = "admission_denied"
                    self._publish("admission_denied", reason=str(payload.get("reason", "denied")))
                    raise AdmissionError("GPU admission denied by local supervisor")
                lease_id = payload.get("lease_id")
                granted = payload.get("granted_vram_bytes")
                if not isinstance(lease_id, str) or not lease_id:
                    raise AdmissionError("GPU admission ACK has no lease identity")
                if self._lease_id is not None and lease_id != self._lease_id:
                    raise AdmissionError("GPU admission ownership changed")
                if isinstance(granted, bool) or not isinstance(granted, int) or granted < 0:
                    raise AdmissionError("GPU admission ACK has an invalid VRAM grant")
                self._lease_id = lease_id
                self._granted_vram_bytes = granted
                return
            waited = (datetime.now(timezone.utc) - started).total_seconds()
            pause = min(self._admission_poll_seconds, remaining)
            deadline.wait(pause)
            remaining -= max(pause, waited)
        self._phase = "admission_denied"
        self._publish("admission_denied", reason="ack_timeout")
        raise AdmissionError("timed out waiting for GPU admission ACK")

    def _valid_ack(self, payload: object, request: dict) -> bool:
        return bool(
            isinstance(payload, dict)
            and payload.get("schema_version") == SCHEMA_VERSION
            and payload.get("run_id") == self._run_id
            and payload.get("request_event_id") == request["event_id"]
            and payload.get("decision") in {"grant", "deny"}
        )

    def _start_heartbeat(self) -> None:
        if not self.enabled or (self._thread is not None and self._thread.is_alive()):
            return
        self._stop.clear()
        self._thread = threading.Thread(
            target=self._heartbeat_loop,
            name="obliteratus-gpu-heartbeat",
            daemon=True,
        )
        self._thread.start()

    def _heartbeat_loop(self) -> None:
        while not self._stop.wait(self._heartbeat_seconds):
            self.heartbeat()

    def _publish(self, event: str, *, reason: str | None = None) -> dict | None:
        if not self.enabled:
            return None
        self._sequence += 1
        payload = {
            "schema_version": SCHEMA_VERSION,
            "event_id": f"{self._run_id}:{self._sequence}",
            "sequence": self._sequence,
            "event": event,
            "phase": self._phase,
            "run_id": self._run_id,
            "model_id": self._model_id,
            "pid": os.getpid(),
            "timestamp": self._clock().isoformat(),
            "allocated_vram_bytes": self._last_memory.allocated_bytes,
            "reserved_vram_bytes": self._last_memory.reserved_bytes,
            "device_count": self._last_memory.device_count,
        }
        if reason is not None:
            payload["reason"] = reason
        encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"))
        events = self._dir / "events.jsonl"  # type: ignore[operator]
        current = self._dir / "current.json"  # type: ignore[operator]
        temporary = self._dir / f".current.{os.getpid()}.tmp"  # type: ignore[operator]
        with events.open("a", encoding="utf-8") as stream:
            stream.write(encoded + "\n")
            stream.flush()
        temporary.write_text(encoded + "\n", encoding="utf-8")
        os.replace(temporary, current)
        return payload


def from_environment() -> GpuLifecyclePublisher:
    """Build the process publisher; an unset directory yields a no-op publisher."""
    interval = os.environ.get("OBLITERATUS_GPU_HEARTBEAT_SECONDS", "15")
    timeout = os.environ.get("OBLITERATUS_GPU_ADMISSION_TIMEOUT_SECONDS", "30")
    try:
        heartbeat_seconds = float(interval)
    except ValueError:
        heartbeat_seconds = 15.0
    try:
        admission_timeout_seconds = float(timeout)
    except ValueError:
        admission_timeout_seconds = 30.0
    return GpuLifecyclePublisher(
        os.environ.get("OBLITERATUS_GPU_LIFECYCLE_DIR"),
        heartbeat_seconds=heartbeat_seconds,
        admission_timeout_seconds=admission_timeout_seconds,
    )


def measure_torch_memory(torch_module) -> MemoryUsage:
    """Measure this process' CUDA allocator without initializing CUDA on CPU hosts."""
    cuda = getattr(torch_module, "cuda", None)
    if cuda is None or not cuda.is_available():
        return MemoryUsage()
    count = cuda.device_count()
    return MemoryUsage(
        allocated_bytes=sum(int(cuda.memory_allocated(index)) for index in range(count)),
        reserved_bytes=sum(int(cuda.memory_reserved(index)) for index in range(count)),
        device_count=count,
    )