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


@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,
        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._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._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 = "loading"
            return self._publish("loading")

    def resize(self, memory: MemoryUsage) -> dict | None:
        with self._lock:
            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()
            return event

    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")
    try:
        heartbeat_seconds = float(interval)
    except ValueError:
        heartbeat_seconds = 15.0
    return GpuLifecyclePublisher(
        os.environ.get("OBLITERATUS_GPU_LIFECYCLE_DIR"),
        heartbeat_seconds=heartbeat_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,
    )