diff --git a/backend/lsdj/gpu_broker.py b/backend/lsdj/gpu_broker.py index 6458386..3132dae 100644 --- a/backend/lsdj/gpu_broker.py +++ b/backend/lsdj/gpu_broker.py @@ -30,6 +30,9 @@ SCHEMA_VERSION = 1 MAX_RECORDS = 32 DEFAULT_POLL_SECONDS = 0.05 +PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 +ERROR_INVALID_PARAMETER = 87 +STILL_ACTIVE = 259 class Priority(enum.IntEnum): @@ -58,20 +61,57 @@ class Lease: pid: int +def _windows_pid_alive(pid: int) -> bool: + """Query one Windows process without treating signal emulation as truth. + + ``os.kill(pid, 0)`` is not a portable liveness probe on Windows. Open the + process with the least query privilege instead, and retain an indeterminate + broker record whenever access is denied or a query fails. False is safe + only when Windows positively reports an invalid PID or a completed process. + """ + + import ctypes + from ctypes import wintypes + + try: + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + open_process = kernel32.OpenProcess + open_process.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + open_process.restype = wintypes.HANDLE + get_exit_code = kernel32.GetExitCodeProcess + get_exit_code.argtypes = [wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD)] + get_exit_code.restype = wintypes.BOOL + close_handle = kernel32.CloseHandle + close_handle.argtypes = [wintypes.HANDLE] + close_handle.restype = wintypes.BOOL + handle = open_process(PROCESS_QUERY_LIMITED_INFORMATION, False, pid) + except (AttributeError, OSError): + return True + if not handle: + return ctypes.get_last_error() != ERROR_INVALID_PARAMETER + exit_code = wintypes.DWORD() + try: + try: + queried = get_exit_code(handle, ctypes.byref(exit_code)) + except OSError: + return True + return not queried or exit_code.value == STILL_ACTIVE + finally: + close_handle(handle) + + def _pid_alive(pid: int) -> bool: if pid <= 0: return False if pid == os.getpid(): return True + if os.name == "nt": + return _windows_pid_alive(pid) try: os.kill(pid, 0) except ProcessLookupError: return False - except PermissionError: - return True - except OSError: - # Windows can reject signal 0 for an otherwise-live process. Retaining - # the record is safer than admitting overlapping GPU work. + except (PermissionError, OSError): return True return True diff --git a/backend/lsdj/mrt2_pytorch.py b/backend/lsdj/mrt2_pytorch.py index 3baf1f1..823869d 100644 --- a/backend/lsdj/mrt2_pytorch.py +++ b/backend/lsdj/mrt2_pytorch.py @@ -7,7 +7,6 @@ from __future__ import annotations -import contextlib import importlib.metadata import math import os @@ -19,7 +18,7 @@ import numpy as np from . import runtime_paths -from .gpu_broker import GpuBroker, Priority +from .gpu_broker import BrokerError, GpuBroker, Lease, Priority from .engine import ( CFG_MUSICCOCA, CFG_NOTES, @@ -55,6 +54,7 @@ ) MAX_SEED = (1 << 63) - 1 +GPU_ADMISSION_TIMEOUT_SECONDS = 30.0 @dataclass(frozen=True) @@ -163,16 +163,6 @@ def __init__( if model not in MODEL_SNAPSHOTS: raise ValueError(f"unknown pinned PyTorch MRT2 model {model!r}") self._selection = selection - self._bindings = bindings or load_bindings() - torch = self._bindings.torch - if not torch.cuda.is_available(): - raise RuntimeUnavailable( - "PyTorch reports no CUDA accelerator; MRT2 has no CPU fallback" - ) - if not getattr(torch.version, "cuda", None): - raise RuntimeUnavailable( - "the installed PyTorch build has no CUDA runtime; MRT2 has no CPU fallback" - ) if cache_root is None: assets = runtime_paths.assets_home() @@ -202,9 +192,45 @@ def __init__( ), ) - # `trust_remote_code` is safe only because model_path resolves the exact - # installer-verified revision above. Never pass a mutable repository ID. + broker_root = runtime_paths.cache_home() + if gpu_broker is None: + if broker_root is None: + raise RuntimeUnavailable( + "LSDJ_CACHE_HOME is missing; MRT2 cannot allocate CUDA without " + "the shared GPU broker" + ) + self._gpu_broker = GpuBroker(broker_root / "gpu-broker") + else: + self._gpu_broker = gpu_broker + self._gpu_lease: Lease | None = None + + # The lease covers the model's full CUDA lifetime, not just generate(). + # Acquire before importing/querying torch: cuda.is_available(), model.to(), + # and processor loading may all initialize a context or reserve VRAM. It + # intentionally lives until worker-process exit, the only reliable CUDA + # context teardown boundary; PID pruning then removes the broker record. try: + self._gpu_lease = self._gpu_broker.acquire( + "mrt2", + priority=Priority.MRT2_REALTIME, + reservation_bytes=0, + capacity_bytes=0, + timeout_seconds=GPU_ADMISSION_TIMEOUT_SECONDS, + ) + self._bindings = bindings or load_bindings() + torch = self._bindings.torch + if not torch.cuda.is_available(): + raise RuntimeUnavailable( + "PyTorch reports no CUDA accelerator; MRT2 has no CPU fallback" + ) + if not getattr(torch.version, "cuda", None): + raise RuntimeUnavailable( + "the installed PyTorch build has no CUDA runtime; " + "MRT2 has no CPU fallback" + ) + + # `trust_remote_code` is safe only because model_path resolves the exact + # installer-verified revision above. Never pass a mutable repository ID. upstream = self._bindings.auto_model.from_pretrained( model_path, trust_remote_code=True, @@ -213,6 +239,8 @@ def __init__( ) self._system = upstream.to("cuda").eval() self._system.load_processor(processor_path, device="cuda") + except (BrokerError, RuntimeUnavailable): + raise except Exception as error: raise RuntimeUnavailable( "the pinned PyTorch MRT2 snapshot could not initialize on CUDA" @@ -221,14 +249,6 @@ def __init__( self._model = model self._model_pin = model_pin self._model_lock = threading.RLock() - broker_root = runtime_paths.cache_home() - self._gpu_broker = ( - gpu_broker - if gpu_broker is not None - else None - if broker_root is None - else GpuBroker(broker_root / "gpu-broker") - ) self._warmup_owner = True self._init_deck_state() @@ -258,6 +278,7 @@ def shared_deck(self) -> "PytorchMrt2Engine": deck._model_pin = self._model_pin deck._model_lock = self._model_lock deck._gpu_broker = self._gpu_broker + deck._gpu_lease = self._gpu_lease deck._warmup_owner = False deck._init_deck_state() return deck @@ -394,40 +415,23 @@ def _generate( ) -> tuple[np.ndarray, Any]: notes = self._notes if stream_conditioning else None drums = self._drums if stream_conditioning else None - broker_hold = ( - contextlib.nullcontext() - if self._gpu_broker is None - else self._gpu_broker.hold( - "mrt2", - priority=Priority.MRT2_REALTIME, - reservation_bytes=0, - capacity_bytes=int( - self._bindings.torch.cuda.get_device_properties( - self._bindings.torch.cuda.current_device() - ).total_memory - ), - timeout_seconds=max(10.0, frames * FRAME_SECONDS), + # The process-lifetime lease was acquired before CUDA initialization. + # The lock only serializes the two deck states sharing this loaded model. + with self._model_lock: + audio, state = self._system.generate( + style=style, + notes=notes, + drums=None if drums is None else [drums], + cfg_drums=self._drums_cfg if stream_conditioning else None, + temperature=self._temperature, + top_k=self._top_k, + cfg_musiccoca=self._cfg_musiccoca, + cfg_notes=self._cfg_notes, + frames=frames, + seed=self._seed, + state=state, + guidance=True, ) - ) - # Acquire the cross-process priority lease before the in-process model - # lock. A waiting MRT2 lease makes a background SA3 callback cancel its - # disposable process, while the two deck states remain serialized here. - with broker_hold: - with self._model_lock: - audio, state = self._system.generate( - style=style, - notes=notes, - drums=None if drums is None else [drums], - cfg_drums=self._drums_cfg if stream_conditioning else None, - temperature=self._temperature, - top_k=self._top_k, - cfg_musiccoca=self._cfg_musiccoca, - cfg_notes=self._cfg_notes, - frames=frames, - seed=self._seed, - state=state, - guidance=True, - ) samples = np.asarray(audio) expected = frames * round(SAMPLE_RATE * FRAME_SECONDS) if samples.ndim != 2 or samples.shape != (expected, CHANNELS): diff --git a/backend/tests/test_gpu_broker.py b/backend/tests/test_gpu_broker.py index c33366f..7dd3147 100644 --- a/backend/tests/test_gpu_broker.py +++ b/backend/tests/test_gpu_broker.py @@ -1,8 +1,10 @@ +import ctypes import json import pathlib import pytest +from lsdj import gpu_broker from lsdj.gpu_broker import ( BrokerCancelled, BrokerError, @@ -12,10 +14,107 @@ ) +class FakeWindowsCall: + def __init__(self, callback): + self.callback = callback + self.argtypes = None + self.restype = None + + def __call__(self, *args): + return self.callback(*args) + + +class FakeKernel32: + def __init__( + self, + *, + handle=41, + exit_code=gpu_broker.STILL_ACTIVE, + query_succeeds=True, + ): + self.handle = handle + self.exit_code = exit_code + self.query_succeeds = query_succeeds + self.closed = [] + self.OpenProcess = FakeWindowsCall(self._open_process) + self.GetExitCodeProcess = FakeWindowsCall(self._get_exit_code) + self.CloseHandle = FakeWindowsCall(self._close_handle) + + def _open_process(self, access, inherit, pid): + assert access == gpu_broker.PROCESS_QUERY_LIMITED_INFORMATION + assert inherit is False + assert pid > 0 + return self.handle + + def _get_exit_code(self, handle, destination): + assert handle == self.handle + destination._obj.value = self.exit_code + return self.query_succeeds + + def _close_handle(self, handle): + self.closed.append(handle) + return True + + +def install_fake_windows(monkeypatch, kernel, *, last_error=0): + monkeypatch.setattr( + ctypes, "WinDLL", lambda *_args, **_kwargs: kernel, raising=False + ) + monkeypatch.setattr(ctypes, "get_last_error", lambda: last_error, raising=False) + + def broker(tmp_path: pathlib.Path) -> GpuBroker: return GpuBroker(tmp_path / "gpu-broker", poll_seconds=0.001) +@pytest.mark.parametrize( + ("exit_code", "expected"), + [(gpu_broker.STILL_ACTIVE, True), (0, False)], +) +def test_windows_liveness_queries_exit_state_and_closes_handle( + monkeypatch, exit_code, expected +): + kernel = FakeKernel32(exit_code=exit_code) + install_fake_windows(monkeypatch, kernel) + + assert gpu_broker._windows_pid_alive(1234) is expected + assert kernel.closed == [kernel.handle] + + +@pytest.mark.parametrize( + ("last_error", "expected"), + [(gpu_broker.ERROR_INVALID_PARAMETER, False), (5, True), (12345, True)], +) +def test_windows_open_failure_only_prunes_a_definitively_invalid_pid( + monkeypatch, last_error, expected +): + kernel = FakeKernel32(handle=0) + install_fake_windows(monkeypatch, kernel, last_error=last_error) + + assert gpu_broker._windows_pid_alive(1234) is expected + assert kernel.closed == [] + + +def test_windows_failed_exit_query_fails_closed_and_closes_handle(monkeypatch): + kernel = FakeKernel32(query_succeeds=False) + install_fake_windows(monkeypatch, kernel) + + assert gpu_broker._windows_pid_alive(1234) is True + assert kernel.closed == [kernel.handle] + + +@pytest.mark.parametrize( + ("error", "expected"), + [(ProcessLookupError(), False), (PermissionError(), True), (OSError(), True)], +) +def test_posix_liveness_semantics_are_preserved(monkeypatch, error, expected): + def fail(_pid, _signal): + raise error + + monkeypatch.setattr(gpu_broker.os, "kill", fail) + assert gpu_broker._pid_alive(gpu_broker.os.getpid() + 1000) is expected + + def test_sa3_lease_is_bounded_by_measured_capacity(tmp_path): service = broker(tmp_path) with pytest.raises(BrokerTimeout): diff --git a/backend/tests/test_mrt2_pytorch.py b/backend/tests/test_mrt2_pytorch.py index fbc9803..0fe1a34 100644 --- a/backend/tests/test_mrt2_pytorch.py +++ b/backend/tests/test_mrt2_pytorch.py @@ -1,7 +1,6 @@ from pathlib import Path from types import SimpleNamespace import tempfile -from contextlib import contextmanager import numpy as np import pytest @@ -13,15 +12,44 @@ RuntimeSelection, RuntimeUnavailable, ) -from lsdj.mrt2_pytorch import PytorchBindings, PytorchMrt2Engine -from lsdj.gpu_broker import Priority +from lsdj.mrt2_pytorch import ( + GPU_ADMISSION_TIMEOUT_SECONDS, + PytorchBindings, + PytorchMrt2Engine, +) +from lsdj.gpu_broker import BrokerCancelled, Priority + + +class RecordingBroker: + def __init__(self, events=None, *, error=None): + self.events = events if events is not None else [] + self.error = error + self.calls = [] + self.releases = [] + self.lease = object() + + def acquire(self, service, **kwargs): + self.events.append("broker_acquire") + self.calls.append((service, kwargs)) + if self.error is not None: + raise self.error + return self.lease + + def release(self, lease): + self.events.append("broker_release") + self.releases.append(lease) + + +DEFAULT_TEST_BROKER = object() class FakeCuda: - def __init__(self, available=True): + def __init__(self, available=True, events=None): self.available = available + self.events = events if events is not None else [] def is_available(self): + self.events.append("cuda_available") return self.available def current_device(self): @@ -37,8 +65,8 @@ def get_device_capability(self, _index): class FakeTorch: bfloat16 = "bf16" - def __init__(self, available=True): - self.cuda = FakeCuda(available) + def __init__(self, available=True, events=None): + self.cuda = FakeCuda(available, events) self.version = SimpleNamespace(cuda="13.0") self._C = SimpleNamespace(_cuda_getDriverVersion=lambda: 13020) @@ -60,20 +88,26 @@ def tokenize(self, embedding): class FakeModel: - def __init__(self): + def __init__(self, events=None, *, fail_cuda_load=False): self.processor = FakeProcessor() self.calls = [] self.processor_path = None self.bad_shape = False + self.events = events if events is not None else [] + self.fail_cuda_load = fail_cuda_load def to(self, device): assert device == "cuda" + self.events.append("model_to_cuda") + if self.fail_cuda_load: + raise RuntimeError("injected CUDA allocation failure") return self def eval(self): return self def load_processor(self, path, *, device): + self.events.append("processor_to_cuda") self.processor_path = (path, device) def generate(self, **kwargs): @@ -86,20 +120,25 @@ def generate(self, **kwargs): class FakeAutoModel: - def __init__(self, model): + def __init__(self, model, events=None): self.model = model self.calls = [] + self.events = events if events is not None else [] def from_pretrained(self, path, **kwargs): + self.events.append("from_pretrained") self.calls.append((path, kwargs)) return self.model -def make_engine(*, cuda=True, gpu_broker=None): - model = FakeModel() - auto_model = FakeAutoModel(model) +def make_engine(*, cuda=True, gpu_broker=DEFAULT_TEST_BROKER, events=None, model=None): + events = events if events is not None else [] + if gpu_broker is DEFAULT_TEST_BROKER: + gpu_broker = RecordingBroker(events) + model = model if model is not None else FakeModel(events) + auto_model = FakeAutoModel(model, events) bindings = PytorchBindings( - torch=FakeTorch(cuda), + torch=FakeTorch(cuda, events), auto_model=auto_model, versions={ "torch": "2.12.1", @@ -149,6 +188,16 @@ def test_cuda_is_mandatory_and_never_falls_back_to_cpu(): make_engine(cuda=False) +def test_missing_broker_root_fails_before_cuda_is_touched(monkeypatch): + events = [] + monkeypatch.delenv("LSDJ_CACHE_HOME", raising=False) + + with pytest.raises(RuntimeUnavailable, match="shared GPU broker"): + make_engine(gpu_broker=None, events=events) + + assert events == [] + + def test_weighted_style_and_controls_map_to_upstream_generate(): engine, model, _, _ = make_engine() engine.set_style([("funk", 3.0), ("dub", 1.0)]) @@ -213,6 +262,7 @@ def test_shared_deck_reuses_one_model_with_independent_continuation_state(): assert first._system is second._system assert first._model_lock is second._model_lock assert first._gpu_broker is second._gpu_broker + assert first._gpu_lease is second._gpu_lease first.generate_chunk() second.generate_chunk() @@ -248,27 +298,64 @@ def test_diagnostics_disclose_unqualified_runtime_and_cuda_versions(): assert diagnostics["capabilities"]["negative_prompt"] is False -def test_mrt2_generation_takes_realtime_priority_over_background_sa3(): - class FakeBroker: - def __init__(self): - self.calls = [] - - @contextmanager - def hold(self, service, **kwargs): - self.calls.append((service, kwargs)) - yield object() - - broker = FakeBroker() +def test_mrt2_model_lifetime_takes_realtime_priority_over_background_sa3(): + broker = RecordingBroker() engine, _, _, _ = make_engine(gpu_broker=broker) engine.generate_chunk() - service, values = broker.calls[-1] + assert len(broker.calls) == 1, "generation must reuse the model-lifetime lease" + service, values = broker.calls[0] assert service == "mrt2" assert values["priority"] is Priority.MRT2_REALTIME assert values["reservation_bytes"] == 0 - assert values["capacity_bytes"] == 12 * 1024**3 + assert values["capacity_bytes"] == 0 + assert values["timeout_seconds"] == GPU_ADMISSION_TIMEOUT_SECONDS + assert broker.releases == [], "only worker-process exit may release CUDA ownership" assert engine.diagnostics()["gpu_broker"] == { "enabled": True, "priority": 100, "preempts": "sa3-background", } + + +def test_gpu_lease_is_acquired_before_any_cuda_or_model_allocation(): + events = [] + broker = RecordingBroker(events) + + make_engine(gpu_broker=broker, events=events) + + assert events == [ + "broker_acquire", + "cuda_available", + "from_pretrained", + "model_to_cuda", + "processor_to_cuda", + ] + + +def test_cancelled_gpu_admission_never_touches_cuda_or_the_model(): + events = [] + broker = RecordingBroker(events, error=BrokerCancelled("injected cancellation")) + + with pytest.raises(BrokerCancelled, match="injected cancellation"): + make_engine(gpu_broker=broker, events=events) + + assert events == ["broker_acquire"] + assert broker.releases == [] + + +def test_cuda_load_failure_keeps_the_lease_until_process_cleanup(): + events = [] + broker = RecordingBroker(events) + model = FakeModel(events, fail_cuda_load=True) + + with pytest.raises(RuntimeUnavailable, match="could not initialize on CUDA"): + make_engine(gpu_broker=broker, events=events, model=model) + + assert events == [ + "broker_acquire", + "cuda_available", + "from_pretrained", + "model_to_cuda", + ] + assert broker.releases == [], "failed CUDA teardown is safe only at process exit" diff --git a/docs/adr/0038-windows-sa3-cuda-qualification-gate.md b/docs/adr/0038-windows-sa3-cuda-qualification-gate.md index 57d7ea4..7c0e647 100644 --- a/docs/adr/0038-windows-sa3-cuda-qualification-gate.md +++ b/docs/adr/0038-windows-sa3-cuda-qualification-gate.md @@ -70,6 +70,40 @@ reservation, provenance, and model path facts before importing a model. A reported free-memory value is only an admission snapshot; the process boundary is still the recovery mechanism for unrelated VRAM pressure and CUDA failure. +## Qualification no-go: resident MRT2 ownership + +Commit `5cbe842` is approved as broker safety behavior. On Windows it prunes a +lease only when the owner is positively known to have exited; access or query +uncertainty retains the record and fails closed. This approval is not NVIDIA +hardware or VRAM evidence. + +Commit `7fdbb32` deliberately acquires one shared, process-lifetime MRT2 lease +before CUDA availability is queried or any model is loaded. Keeping that lease +through model-load failure and for the lifetime of the worker is fail-closed: +SA3 cannot treat VRAM owned by a resident MRT2 model as free. The ordering must +remain until a measured replacement is available. + +The current single-level lease consequently prevents SA3 admission while the +MRT2 worker is alive, so it cannot satisfy the coexistence and dual-deck tests +required by issue #114. Releasing the lease between MRT2 generations would not +be safe because the model remains resident. It would also undercount ownership +when multiple models or workers are present. + +Issue #114 is a no-go until one of these paths is completed: + +1. Measure resident and active-generation VRAM on Windows NVIDIA hardware, then + implement a two-level broker. Long-lived reservations must account for every + resident model and worker; short active-generation priority leases must let + realtime MRT2 work interrupt or defer background SA3 work. +2. Obtain explicit product acceptance that MRT2 will unload before SA3 starts, + then re-load and warm up afterward, and qualify that lifecycle instead of + coexistence. + +The required VRAM and hardware evidence does not yet exist. Auto selection, +public SA3 CUDA availability, `HARDWARE_QUALIFIED`, and the release manifest +gate remain disabled. This record does not change runtime code, dependency +pins, qualification gates, or make a release-support claim. + ## Consequences The design and model-free failure behavior can merge without delaying the diff --git a/docs/issue-114-windows-sa3-cuda-checklist.md b/docs/issue-114-windows-sa3-cuda-checklist.md index 80096bf..7b396c4 100644 --- a/docs/issue-114-windows-sa3-cuda-checklist.md +++ b/docs/issue-114-windows-sa3-cuda-checklist.md @@ -9,6 +9,39 @@ Set `LSDJ_ALLOW_UNVERIFIED_SA3_CUDA=1` only on a dedicated qualification host. It permits explicit GPU probes; it does not enable Auto or make a build release-ready. +## Current no-go: resident MRT2 ownership + +The following checks record reviewed code behavior, not Windows NVIDIA +qualification evidence: + +- [x] `5cbe842` retains Windows lease records when owner-process liveness is + uncertain and prunes only owners positively known to have exited. This + fail-closed safety behavior is approved. +- [x] `7fdbb32` acquires the shared MRT2 lease before CUDA or model load and + holds it for the worker lifetime, including after a CUDA load failure. This + ordering deliberately prevents resident MRT2 VRAM from being treated as + free. + +**NO-GO:** the process-lifetime MRT2 lease prevents SA3 admission while the +MRT2 worker is alive. Releasing it between generations is unsafe because model +VRAM remains resident. Issue #114 cannot pass its broker or dual-deck +acceptance sections until one of these alternatives is complete: + +- [ ] On Windows NVIDIA hardware, measure resident and active-generation VRAM + separately for each MRT2 model and worker, both decks/multiple model + combinations, and SA3 Small Music and Small SFX. +- [ ] Implement and validate a two-level broker: long-lived reservations for + every resident model/worker plus short active-generation priority leases for + realtime MRT2 versus background SA3 work. Prove that capacity is neither + double-counted nor inferred from VRAM that remains resident. +- [ ] Alternatively, obtain explicit product acceptance to unload MRT2 before + SA3 starts and re-load/warm it afterward, then replace the coexistence tests + below with evidence for that lifecycle. + +Until that decision and the required hardware evidence exist, Auto and public +SA3 CUDA remain disabled, `HARDWARE_QUALIFIED` and the release gate remain +unchanged, and no Windows SA3 CUDA release claim may be made. + ## Immutable inputs and shared runtime - [x] Pin the official upstream source commit without an LSDJ fork. @@ -86,8 +119,9 @@ route it explicitly to TFLite. - [ ] Start SA3, then request MRT2 work during loading, sampling, and decoding. The watchdog/callback exits SA3, releases its lease/context, and MRT2 proceeds. -- [ ] Queue SA3 while MRT2 holds a lease. SA3 waits without disturbing either - deck or the native audio callback. +- [ ] With the two-level broker, queue SA3 while MRT2 holds an active-generation + priority lease. SA3 waits without disturbing either deck or the native audio + callback; MRT2's resident reservation remains accounted for. - [ ] Cancel while waiting, loading, sampling, and decoding; no child or CUDA allocation remains. - [ ] Force CUDA OOM, worker exception, invalid output, and abrupt worker death;