mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 14:40:38 +02:00
356 lines
14 KiB
Python
356 lines
14 KiB
Python
"""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 logging
|
|
import os
|
|
from pathlib import Path
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from typing import Callable
|
|
|
|
|
|
SCHEMA_VERSION = 1
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AdmissionError(RuntimeError):
|
|
"""The local supervisor did not grant correlated GPU admission."""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str,
|
|
*,
|
|
reason: str = "admission_failed",
|
|
run_id: str | None = None,
|
|
request_event_id: str | None = None,
|
|
elapsed_seconds: float | None = None,
|
|
timeout_seconds: float | None = None,
|
|
) -> None:
|
|
super().__init__(message)
|
|
self.reason = reason
|
|
self.run_id = run_id
|
|
self.request_event_id = request_event_id
|
|
self.elapsed_seconds = elapsed_seconds
|
|
self.timeout_seconds = timeout_seconds
|
|
|
|
def diagnostic_message(self) -> str:
|
|
"""Return operator-safe correlated details without lease credentials."""
|
|
fields = [f"reason={self.reason}"]
|
|
if self.run_id:
|
|
fields.append(f"run_id={self.run_id}")
|
|
if self.request_event_id:
|
|
fields.append(f"request_event_id={self.request_event_id}")
|
|
if self.elapsed_seconds is not None:
|
|
fields.append(f"elapsed={self.elapsed_seconds:.1f}s")
|
|
if self.timeout_seconds is not None:
|
|
fields.append(f"deadline={self.timeout_seconds:.1f}s")
|
|
return ", ".join(fields)
|
|
|
|
def user_message(self) -> str:
|
|
"""Return an actionable UI message for a safe admission failure."""
|
|
if self.reason == "ack_timeout":
|
|
return (
|
|
"GPU admission timed out before the host supervisor responded. "
|
|
"No model weights were loaded or modified. The GPU broker may still "
|
|
"be draining another workload; retry after it becomes available."
|
|
)
|
|
if self.reason.startswith("broker_admission_"):
|
|
return (
|
|
"The host GPU supervisor denied this run because it could not reserve "
|
|
"the requested GPU capacity. No model weights were loaded or modified."
|
|
)
|
|
return f"GPU admission failed: {self}"
|
|
|
|
|
|
@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"
|
|
event = self._publish("allocation_started")
|
|
self._start_heartbeat()
|
|
return event
|
|
|
|
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 in {"released", "intent_published"}:
|
|
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
|
|
started_at = time.monotonic()
|
|
logger.info(
|
|
"GPU admission requested run_id=%s request_event_id=%s timeout=%.1fs",
|
|
self._run_id,
|
|
request["event_id"],
|
|
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":
|
|
reason = str(payload.get("reason", "denied"))
|
|
self._phase = "admission_denied"
|
|
self._publish("admission_denied", reason=reason)
|
|
elapsed = time.monotonic() - started_at
|
|
logger.warning(
|
|
"GPU admission denied run_id=%s request_event_id=%s "
|
|
"reason=%s elapsed=%.1fs",
|
|
self._run_id,
|
|
request["event_id"],
|
|
reason,
|
|
elapsed,
|
|
)
|
|
raise AdmissionError(
|
|
"GPU admission denied by local supervisor",
|
|
reason=reason,
|
|
run_id=self._run_id,
|
|
request_event_id=request["event_id"],
|
|
elapsed_seconds=elapsed,
|
|
timeout_seconds=self._admission_timeout_seconds,
|
|
)
|
|
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",
|
|
reason="invalid_lease_identity",
|
|
run_id=self._run_id,
|
|
request_event_id=request["event_id"],
|
|
elapsed_seconds=time.monotonic() - started_at,
|
|
timeout_seconds=self._admission_timeout_seconds,
|
|
)
|
|
if self._lease_id is not None and lease_id != self._lease_id:
|
|
raise AdmissionError(
|
|
"GPU admission ownership changed",
|
|
reason="lease_ownership_changed",
|
|
run_id=self._run_id,
|
|
request_event_id=request["event_id"],
|
|
elapsed_seconds=time.monotonic() - started_at,
|
|
timeout_seconds=self._admission_timeout_seconds,
|
|
)
|
|
if isinstance(granted, bool) or not isinstance(granted, int) or granted < 0:
|
|
raise AdmissionError(
|
|
"GPU admission ACK has an invalid VRAM grant",
|
|
reason="invalid_vram_grant",
|
|
run_id=self._run_id,
|
|
request_event_id=request["event_id"],
|
|
elapsed_seconds=time.monotonic() - started_at,
|
|
timeout_seconds=self._admission_timeout_seconds,
|
|
)
|
|
self._lease_id = lease_id
|
|
self._granted_vram_bytes = granted
|
|
logger.info(
|
|
"GPU admission granted run_id=%s request_event_id=%s "
|
|
"elapsed=%.1fs granted_vram_bytes=%d",
|
|
self._run_id,
|
|
request["event_id"],
|
|
time.monotonic() - started_at,
|
|
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")
|
|
elapsed = time.monotonic() - started_at
|
|
error = AdmissionError(
|
|
"timed out waiting for GPU admission ACK",
|
|
reason="ack_timeout",
|
|
run_id=self._run_id,
|
|
request_event_id=request["event_id"],
|
|
elapsed_seconds=elapsed,
|
|
timeout_seconds=self._admission_timeout_seconds,
|
|
)
|
|
logger.error("GPU admission timeout %s", error.diagnostic_message())
|
|
raise error
|
|
|
|
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,
|
|
run_id=os.environ.get("OBLITERATUS_RUN_ID"),
|
|
)
|
|
|
|
|
|
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,
|
|
)
|