Harden GPU lease ownership and record SA3 CUDA no-go

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
 

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):

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):