From 4b677be612bfbf6b9f7c43896fff5413704d7af5 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:27:21 -0400 Subject: [PATCH] feat: publish GPU model lifecycle events --- app.py | 63 +++++++---- ci/test-risk-map.json | 11 ++ docs/deployment/shared-gpu-host.md | 35 +++++- obliteratus/gpu_lifecycle.py | 169 +++++++++++++++++++++++++++++ tests/test_gpu_lifecycle.py | 75 +++++++++++++ 5 files changed, 331 insertions(+), 22 deletions(-) create mode 100644 obliteratus/gpu_lifecycle.py create mode 100644 tests/test_gpu_lifecycle.py diff --git a/app.py b/app.py index a60a1bc..56795a2 100644 --- a/app.py +++ b/app.py @@ -58,6 +58,7 @@ if "HF_HOME" not in os.environ: import gradio as gr import torch from obliteratus import device as dev +from obliteratus.gpu_lifecycle import from_environment, measure_torch_memory from obliteratus.credential_sources import resolve_first, resolve_secret, secret_available from obliteratus.ui_vram import ( DEFAULT_VRAM_REFRESH_INTERVAL, @@ -74,6 +75,8 @@ from obliteratus.model_load_settings import ( from obliteratus.persistence_contracts import validate_reloadable_checkpoint from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer +_gpu_lifecycle = from_environment() + # ── ZeroGPU support ───────────────────────────────────────────────── # When running on HuggingFace Spaces with ZeroGPU, the `spaces` package # provides the @spaces.GPU decorator that allocates a GPU from the shared @@ -131,6 +134,7 @@ def _load_model_to_device( backends and explicitly moves the model to the best device after loading. On CUDA the behaviour is identical to ``device_map="auto"``. """ + _gpu_lifecycle.loading(str(pretrained_path)) kwargs: dict = {} if torch_dtype is not None: kwargs["torch_dtype"] = torch_dtype @@ -150,12 +154,20 @@ def _load_model_to_device( if dev.supports_device_map_auto(): kwargs["device_map"] = "auto" - model = AutoModelForCausalLM.from_pretrained(pretrained_path, **kwargs) + try: + model = AutoModelForCausalLM.from_pretrained(pretrained_path, **kwargs) # On MPS / CPU: model loaded without device_map, move to best device - if not dev.supports_device_map_auto(): - target = dev.get_device() - model = model.to(target) + if not dev.supports_device_map_auto(): + target = dev.get_device() + model = model.to(target) + except BaseException: + _gpu_lifecycle.release(reason="load_failed") + raise + + memory = measure_torch_memory(torch) + _gpu_lifecycle.resize(memory) + _gpu_lifecycle.ready(memory) return model @@ -915,12 +927,14 @@ def _resolve_ui_load_settings( # Obliteration # --------------------------------------------------------------------------- -def _clear_gpu(): +def _clear_gpu(*, release_lifecycle: bool = True): """Unload model tensors while retaining the checkpoint for lazy reload.""" with _lock: _state["model"] = None _state["tokenizer"] = None dev.free_gpu_memory() + if release_lifecycle: + _gpu_lifecycle.release(reason="unload") def _checkpoint_is_available(checkpoint: str | None) -> bool: @@ -2292,6 +2306,7 @@ def obliterate(model_choice: str, method_choice: str, log_lines.append(f"Resolved load mode: {load_settings.summary}") log_lines.append("") + _gpu_lifecycle.loading(model_id) worker = threading.Thread(target=run_pipeline, daemon=True) worker.start() @@ -2315,6 +2330,7 @@ def obliterate(model_choice: str, method_choice: str, time.sleep(0.5) if timed_out: + _gpu_lifecycle.release(reason="timeout") terminated = _cancel_pipeline_worker( worker, cancellation, @@ -2349,6 +2365,7 @@ def obliterate(model_choice: str, method_choice: str, # Handle error if error_ref[0] is not None: + _gpu_lifecycle.release(reason="load_failed") _cleanup_failed_pipeline(pipeline_ref[0], save_dir) with _lock: _state["status"] = "idle" @@ -2362,6 +2379,9 @@ def obliterate(model_choice: str, method_choice: str, # Wrapped in try/except to ensure status is never stuck on "obliterating". try: pipeline = pipeline_ref[0] + memory = measure_torch_memory(torch) + _gpu_lifecycle.resize(memory) + _gpu_lifecycle.ready(memory) save_dir = str(_resolve_local_checkpoint(save_dir)) can_generate = pipeline._quality_metrics.get("coherence") is not None @@ -2459,7 +2479,7 @@ def obliterate(model_choice: str, method_choice: str, # Free the float16 model pipeline.handle.model = None pipeline.handle.tokenizer = None - _clear_gpu() + _clear_gpu(release_lifecycle=False) # -- Attempt 1: bitsandbytes 4-bit quantization (fast, memory-efficient) bnb_available = False @@ -2512,7 +2532,7 @@ def obliterate(model_choice: str, method_choice: str, log_lines.append( f"4-bit reload failed: {_format_checkpoint_reload_error(e)}" ) - _clear_gpu() + _clear_gpu(release_lifecycle=False) # -- Attempt 2: CPU offloading (slower but no extra dependencies) if not can_generate: @@ -2576,6 +2596,7 @@ def obliterate(model_choice: str, method_choice: str, if can_generate: log_lines.append(f"LIBERATION COMPLETE in {_elapsed()} \u2014 switch to the Chat tab!") else: + _gpu_lifecycle.release(reason="no_active_model") log_lines.append(f"LIBERATION COMPLETE in {_elapsed()} \u2014 model saved!") log_lines.append("=" * 50) @@ -2605,6 +2626,7 @@ def obliterate(model_choice: str, method_choice: str, except Exception as e: # Ensure status never gets stuck on "obliterating" + _gpu_lifecycle.release(reason="post_load_failed") _cleanup_failed_pipeline(pipeline_ref[0], save_dir) with _lock: _state["status"] = "idle" @@ -5719,18 +5741,21 @@ def launch( Called by ``python app.py`` (HF Spaces) or ``obliteratus ui`` (local). """ - demo.launch( - server_name=server_name, - server_port=server_port, - share=share, - inbrowser=inbrowser, - auth=auth, - css=CSS, - max_threads=max_threads, - js=_JS, - quiet=quiet, - theme=THEME, - ) + try: + demo.launch( + server_name=server_name, + server_port=server_port, + share=share, + inbrowser=inbrowser, + auth=auth, + css=CSS, + max_threads=max_threads, + js=_JS, + quiet=quiet, + theme=THEME, + ) + finally: + _gpu_lifecycle.release(reason="process_shutdown") if __name__ == "__main__": diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index f08290b..c0fdd1f 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -176,6 +176,7 @@ "obliteratus/cli.py", "obliteratus/interactive.py", "obliteratus/local_ui.py", + "obliteratus/gpu_lifecycle.py", "obliteratus/models_client.py", "obliteratus/remote.py", "obliteratus/remote_contracts.py", @@ -743,6 +744,16 @@ ], "conditional_gates": [] }, + { + "path": "obliteratus/gpu_lifecycle.py", + "risk_class": "cpu-contract", + "risk": "local GPU lifecycle event ordering, recovery state, and heartbeat shutdown", + "contract_owner": "operator interface maintainers", + "required_tests": [ + "tests/test_gpu_lifecycle.py" + ], + "conditional_gates": [] + }, { "path": "obliteratus/service_contracts.py", "risk_class": "cpu-contract", diff --git a/docs/deployment/shared-gpu-host.md b/docs/deployment/shared-gpu-host.md index 90e98ff..2c0e202 100644 --- a/docs/deployment/shared-gpu-host.md +++ b/docs/deployment/shared-gpu-host.md @@ -50,9 +50,38 @@ the requested devices and VRAM before OBLITERATUS initializes CUDA, maintain the lease for the entire process, and release it after CUDA allocations terminate. Schedulers that drain another inference service commonly distinguish `acquire` -from `ready`. OBLITERATUS does not yet emit a portable post-model-load readiness -signal. Such deployments must not mark a lease ready immediately after process -start; add and test an application readiness hook first. +from `ready`. Set `OBLITERATUS_GPU_LIFECYCLE_DIR` to an existing runtime +directory writable by the OBLITERATUS service account to enable the local +lifecycle protocol. When it is unset, publication is a no-op and conservative +wrappers should retain their lease for the full process lifetime. + +The application atomically replaces `current.json` and appends ordered JSON +objects to `events.jsonl`. Schema version 1 publishes `loading`, `resize`, +`ready`, `heartbeat`, and `release`, with run/model identity, a monotonic +sequence, process ID, timestamp, and measured allocator bytes. Readiness occurs +after model allocation, never at UI startup. Set +`OBLITERATUS_GPU_HEARTBEAT_SECONDS` to change the default 15-second heartbeat. + +A minimal systemd setup lets the supervisor create the local boundary without +giving the application scheduling authority: + +```ini +[Service] +User=obliteratus +Group=obliteratus-gpu +RuntimeDirectory=obliteratus-gpu-lifecycle +RuntimeDirectoryMode=0770 +Environment=OBLITERATUS_GPU_LIFECYCLE_DIR=/run/obliteratus-gpu-lifecycle +Environment=OBLITERATUS_GPU_HEARTBEAT_SECONDS=15 +ExecStart=/srv/obliteratus/current/.venv/bin/obliteratus ui --host 127.0.0.1 +``` + +The host supervisor should acquire capacity before starting this unit, consume +the runtime files as untrusted structured data, deduplicate by `event_id`, and +release or resize its reservation only after the corresponding application +event. It should treat a stale heartbeat or process exit as a failed lease and +must not pass its control socket, administrative API, or scheduling permissions +to the application account. The requested VRAM must include weights, activation collection, verification, checkpoint snapshots, CUDA context, and a safety margin. Multi-GPU sharding is diff --git a/obliteratus/gpu_lifecycle.py b/obliteratus/gpu_lifecycle.py new file mode 100644 index 0000000..c72de6f --- /dev/null +++ b/obliteratus/gpu_lifecycle.py @@ -0,0 +1,169 @@ +"""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, + ) diff --git a/tests/test_gpu_lifecycle.py b/tests/test_gpu_lifecycle.py new file mode 100644 index 0000000..fe85f7a --- /dev/null +++ b/tests/test_gpu_lifecycle.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from datetime import datetime, timezone +import json +import time + +import pytest + +from obliteratus.gpu_lifecycle import ( + GpuLifecyclePublisher, + MemoryUsage, + measure_torch_memory, +) + + +def test_fake_supervisor_observes_order_identity_and_recoverable_state(tmp_path): + publisher = GpuLifecyclePublisher( + tmp_path, + heartbeat_seconds=0.01, + clock=lambda: datetime(2026, 1, 2, tzinfo=timezone.utc), + run_id="run-1", + ) + memory = MemoryUsage(allocated_bytes=10, reserved_bytes=12, device_count=1) + publisher.loading("org/model") + publisher.resize(memory) + publisher.ready(memory) + time.sleep(1.05) + publisher.release(reason="test_complete") + assert publisher.release(reason="duplicate") is None + + events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()] + assert [event["event"] for event in events] == [ + "loading", "resize", "ready", "heartbeat", "release", + ] + assert [event["sequence"] for event in events] == list(range(1, 6)) + assert len({event["event_id"] for event in events}) == 5 + assert all(event["run_id"] == "run-1" for event in events) + assert events[1]["reserved_vram_bytes"] == 12 + current = json.loads((tmp_path / "current.json").read_text()) + assert current["event"] == "release" + assert current["reason"] == "test_complete" + + +def test_disabled_publisher_is_noop(): + publisher = GpuLifecyclePublisher(None) + assert publisher.loading("model") is None + assert publisher.resize(MemoryUsage()) is None + assert publisher.ready() is None + assert publisher.heartbeat() is None + assert publisher.release() is None + + +def test_runtime_directory_must_exist(tmp_path): + with pytest.raises(ValueError, match="must already exist"): + GpuLifecyclePublisher(tmp_path / "missing") + + +def test_runtime_directory_rejects_symlink(tmp_path): + real = tmp_path / "real" + real.mkdir() + link = tmp_path / "link" + link.symlink_to(real, target_is_directory=True) + with pytest.raises(ValueError, match="must not be a symlink"): + GpuLifecyclePublisher(link) + + +def test_measure_torch_memory_aggregates_devices(): + class FakeCuda: + is_available = staticmethod(lambda: True) + device_count = staticmethod(lambda: 2) + memory_allocated = staticmethod(lambda index: (index + 1) * 10) + memory_reserved = staticmethod(lambda index: (index + 1) * 20) + + usage = measure_torch_memory(type("Torch", (), {"cuda": FakeCuda})()) + assert usage == MemoryUsage(allocated_bytes=30, reserved_bytes=60, device_count=2)