diff --git a/app.py b/app.py index 56795a2..c205b50 100644 --- a/app.py +++ b/app.py @@ -2152,7 +2152,6 @@ def obliterate(model_choice: str, method_choice: str, use_custom = custom_harmful and custom_harmful.strip() dataset_key = get_source_key_from_label(dataset_source_choice) if dataset_source_choice else "builtin" - _clear_gpu() with _lock: if _state["status"] == "obliterating": yield "**Error:** An obliteration is already in progress.", "", gr.update(), gr.update(), gr.update(), gr.update() @@ -2162,6 +2161,11 @@ def obliterate(model_choice: str, method_choice: str, _state["model_name"] = model_choice _state["method"] = method + # Admission must be granted before cleanup can touch the CUDA runtime or + # the worker can enter any model-loading path. + _gpu_lifecycle.loading(model_id) + _clear_gpu(release_lifecycle=False) + with _lock: global _obliterate_counter _obliterate_counter += 1 @@ -2306,7 +2310,6 @@ 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() @@ -2704,7 +2707,11 @@ def chat_respond(message: str, history: list[dict], system_prompt: str, if model_dev.type == "meta": _needs_reload = True elif dev.is_gpu_available() and model_dev.type not in ("cuda", "mps"): + _gpu_lifecycle.loading(_state.get("model_name") or "chat-model") model.to(dev.get_device()) + memory = measure_torch_memory(torch) + _gpu_lifecycle.resize(memory) + _gpu_lifecycle.ready(memory) except Exception: _needs_reload = True @@ -3183,7 +3190,11 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[ if model_dev.type == "meta": _needs_reload = True elif dev.is_gpu_available() and model_dev.type not in ("cuda", "mps"): + _gpu_lifecycle.loading(model_name or "comparison-model") abliterated_model.to(dev.get_device()) + memory = measure_torch_memory(torch) + _gpu_lifecycle.resize(memory) + _gpu_lifecycle.ready(memory) except Exception: _needs_reload = True diff --git a/docs/deployment/shared-gpu-host.md b/docs/deployment/shared-gpu-host.md index 2c0e202..3ff2b5f 100644 --- a/docs/deployment/shared-gpu-host.md +++ b/docs/deployment/shared-gpu-host.md @@ -56,11 +56,40 @@ 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`, +objects to `events.jsonl`. Before any accelerator allocation it publishes a +`loading` request in the `intent_published` phase and waits for a correlated +`ack.json`. Only a valid grant advances through `admission_granted` to +`allocation_started`; denial, timeout, a stale run/request ID, changed lease +identity, or an invalid grant fails closed. Subsequent events include `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. +sequence, process ID, timestamp, and measured allocator bytes. + +The supervisor acknowledgment schema is: + +```json +{ + "schema_version": 1, + "run_id": "copied from the loading request", + "request_event_id": "copied from the loading request", + "decision": "grant", + "lease_id": "stable supervisor-owned lease identity", + "granted_vram_bytes": 137438953472 +} +``` + +For denial, set `decision` to `deny` and optionally include a non-sensitive +`reason`. Write the acknowledgment to a temporary file and atomically rename it +to `ack.json`; never update it in place. The application accepts only an exact +run and request correlation. A model reload or CPU-to-GPU move repeats the +barrier and must retain the same `lease_id`. The grant is a reservation ceiling: +measured growth beyond `granted_vram_bytes` fails the lifecycle contract rather +than becoming ready. Memory decreases are published after release and need no +new admission. + +Readiness occurs after model allocation, never at UI startup. Set +`OBLITERATUS_GPU_HEARTBEAT_SECONDS` to change the default 15-second heartbeat, +and `OBLITERATUS_GPU_ADMISSION_TIMEOUT_SECONDS` to change the default 30-second +fail-closed acknowledgment timeout. A minimal systemd setup lets the supervisor create the local boundary without giving the application scheduling authority: @@ -73,15 +102,19 @@ RuntimeDirectory=obliteratus-gpu-lifecycle RuntimeDirectoryMode=0770 Environment=OBLITERATUS_GPU_LIFECYCLE_DIR=/run/obliteratus-gpu-lifecycle Environment=OBLITERATUS_GPU_HEARTBEAT_SECONDS=15 +Environment=OBLITERATUS_GPU_ADMISSION_TIMEOUT_SECONDS=30 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 runtime directory must be writable by the application group and readable by +the root supervisor; `ack.json` should be root-owned and group-readable. The +supervisor watches `current.json`, completes acquire/prepare, atomically writes +the correlated acknowledgment, and only then permits the application wait to +finish. On supervisor restart it must recover or explicitly deny the current +intent before writing a new acknowledgment. It should consume runtime files as +untrusted structured data, deduplicate by `event_id`, treat a stale heartbeat or +process exit as a failed lease, and never pass its control 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 index c72de6f..e6b3bcd 100644 --- a/obliteratus/gpu_lifecycle.py +++ b/obliteratus/gpu_lifecycle.py @@ -16,6 +16,10 @@ 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.""" @@ -33,6 +37,8 @@ class GpuLifecyclePublisher: 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: @@ -41,6 +47,8 @@ class GpuLifecyclePublisher: 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() @@ -48,6 +56,8 @@ class GpuLifecyclePublisher: 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: @@ -62,11 +72,25 @@ class GpuLifecyclePublisher: def loading(self, model_id: str) -> dict | None: with self._lock: self._model_id = str(model_id) - self._phase = "loading" - return self._publish("loading") + 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") @@ -94,8 +118,53 @@ class GpuLifecyclePublisher: 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 @@ -146,13 +215,19 @@ class GpuLifecyclePublisher: 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, ) diff --git a/tests/test_gpu_lifecycle.py b/tests/test_gpu_lifecycle.py index fe85f7a..8a5343b 100644 --- a/tests/test_gpu_lifecycle.py +++ b/tests/test_gpu_lifecycle.py @@ -2,11 +2,13 @@ from __future__ import annotations from datetime import datetime, timezone import json +import threading import time import pytest from obliteratus.gpu_lifecycle import ( + AdmissionError, GpuLifecyclePublisher, MemoryUsage, measure_torch_memory, @@ -20,6 +22,19 @@ def test_fake_supervisor_observes_order_identity_and_recoverable_state(tmp_path) clock=lambda: datetime(2026, 1, 2, tzinfo=timezone.utc), run_id="run-1", ) + def acknowledge(): + while not (tmp_path / "current.json").exists(): + time.sleep(0.001) + request = json.loads((tmp_path / "current.json").read_text()) + (tmp_path / "ack.json").write_text(json.dumps({ + "schema_version": 1, + "run_id": "run-1", + "request_event_id": request["event_id"], + "decision": "grant", + "lease_id": "lease-1", + "granted_vram_bytes": 100, + })) + threading.Thread(target=acknowledge).start() memory = MemoryUsage(allocated_bytes=10, reserved_bytes=12, device_count=1) publisher.loading("org/model") publisher.resize(memory) @@ -30,12 +45,13 @@ def test_fake_supervisor_observes_order_identity_and_recoverable_state(tmp_path) 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", + "loading", "admission_granted", "allocation_started", "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 [event["sequence"] for event in events] == list(range(1, 8)) + assert len({event["event_id"] for event in events}) == 7 assert all(event["run_id"] == "run-1" for event in events) - assert events[1]["reserved_vram_bytes"] == 12 + assert events[3]["reserved_vram_bytes"] == 12 current = json.loads((tmp_path / "current.json").read_text()) assert current["event"] == "release" assert current["reason"] == "test_complete" @@ -73,3 +89,96 @@ def test_measure_torch_memory_aggregates_devices(): usage = measure_torch_memory(type("Torch", (), {"cuda": FakeCuda})()) assert usage == MemoryUsage(allocated_bytes=30, reserved_bytes=60, device_count=2) + + +@pytest.mark.parametrize("decision", ["deny", "timeout", "stale"]) +def test_admission_failures_never_enter_allocation(tmp_path, decision): + publisher = GpuLifecyclePublisher( + tmp_path, + admission_timeout_seconds=0.05, + admission_poll_seconds=0.005, + run_id="current-run", + ) + allocation_entered = False + + def acknowledge(): + nonlocal allocation_entered + while not (tmp_path / "current.json").exists(): + time.sleep(0.001) + request = json.loads((tmp_path / "current.json").read_text()) + if decision != "timeout": + (tmp_path / "ack.json").write_text(json.dumps({ + "schema_version": 1, + "run_id": "stale-run" if decision == "stale" else "current-run", + "request_event_id": request["event_id"], + "decision": "deny" if decision == "deny" else "grant", + "lease_id": "lease-1", + "granted_vram_bytes": 100, + })) + + worker = threading.Thread(target=acknowledge) + worker.start() + with pytest.raises(AdmissionError): + publisher.loading("org/model") + allocation_entered = True + worker.join() + assert allocation_entered is False + events = [json.loads(line)["event"] for line in (tmp_path / "events.jsonl").read_text().splitlines()] + assert "allocation_started" not in events + + +def test_delayed_ack_blocks_until_granted(tmp_path): + publisher = GpuLifecyclePublisher( + tmp_path, admission_timeout_seconds=1, admission_poll_seconds=0.005, run_id="run", + ) + entered = threading.Event() + + def load(): + publisher.loading("model") + entered.set() + + worker = threading.Thread(target=load) + worker.start() + time.sleep(0.03) + assert not entered.is_set() + request = json.loads((tmp_path / "current.json").read_text()) + (tmp_path / "ack.json").write_text(json.dumps({ + "schema_version": 1, "run_id": "run", "request_event_id": request["event_id"], + "decision": "grant", "lease_id": "lease", "granted_vram_bytes": 1, + })) + worker.join(timeout=1) + assert entered.is_set() + + +def test_grant_ceiling_and_lease_identity_fail_closed(tmp_path): + publisher = GpuLifecyclePublisher( + tmp_path, admission_timeout_seconds=0.2, admission_poll_seconds=0.002, run_id="run", + ) + + def grant(request_id, lease_id="lease-a", granted=20): + (tmp_path / "ack.json").write_text(json.dumps({ + "schema_version": 1, "run_id": "run", "request_event_id": request_id, + "decision": "grant", "lease_id": lease_id, "granted_vram_bytes": granted, + })) + + def acknowledge_first(): + while not (tmp_path / "current.json").exists(): + time.sleep(0.001) + grant(json.loads((tmp_path / "current.json").read_text())["event_id"]) + + threading.Thread(target=acknowledge_first).start() + publisher.loading("first") + with pytest.raises(AdmissionError, match="exceeds"): + publisher.resize(MemoryUsage(reserved_bytes=21)) + + def acknowledge_changed_owner(): + seen = None + while seen != "loading": + current = json.loads((tmp_path / "current.json").read_text()) + seen = current["event"] + time.sleep(0.001) + grant(current["event_id"], lease_id="lease-b") + + threading.Thread(target=acknowledge_changed_owner).start() + with pytest.raises(AdmissionError, match="ownership changed"): + publisher.loading("second")