feat(registry): active-model swap with async lock and SSE event bus

server/registry.py

@@ -0,0 +1,98 @@
+"""Active-model registry with async swap lock and SSE event bus."""
+from __future__ import annotations
+
+import asyncio
+import gc
+from typing import AsyncIterator
+
+import torch
+
+from server.models.base import ModelAdapter
+
+
+class Registry:
+    def __init__(self, adapter_classes: dict[str, type], device: str):
+        self._classes = adapter_classes
+        self._device = device
+        self._active: ModelAdapter | None = None
+        self._active_id: str | None = None
+        self._status: str = "idle"
+        self._last_error: str | None = None
+        self._lock = asyncio.Lock()
+        self._subscribers: list[asyncio.Queue[dict]] = []
+
+    @property
+    def device(self) -> str:
+        return self._device
+
+    def status(self) -> dict:
+        return {"id": self._active_id, "status": self._status, "last_error": self._last_error}
+
+    def list_models(self) -> list[dict]:
+        return [
+            {
+                "id": cls.id,
+                "label": cls.label,
+                "description": cls.description,
+                "languages": [l.model_dump() for l in cls.languages],
+                "paralinguistic_tags": cls.paralinguistic_tags,
+                "supports_voice_clone": cls.supports_voice_clone,
+                "params": [p.model_dump() for p in cls.params],
+            }
+            for cls in self._classes.values()
+        ]
+
+    async def get_or_load(self, model_id: str) -> ModelAdapter:
+        if model_id not in self._classes:
+            raise KeyError(model_id)
+        async with self._lock:
+            if self._active_id == model_id and self._active is not None:
+                return self._active
+            await self._publish({"id": model_id, "status": "loading"})
+            self._status = "loading"
+            self._last_error = None
+            if self._active is not None:
+                try:
+                    self._active.unload()
+                finally:
+                    self._active = None
+                    self._free_caches()
+            try:
+                instance = self._classes[model_id](self._device)
+                instance.load()
+            except Exception as exc:
+                self._status = "error"
+                self._last_error = str(exc)
+                self._active_id = None
+                await self._publish({"id": model_id, "status": "error", "error": str(exc)})
+                raise
+            self._active = instance
+            self._active_id = model_id
+            self._status = "loaded"
+            await self._publish({"id": model_id, "status": "loaded"})
+            return instance
+
+    async def stream_events(self) -> AsyncIterator[dict]:
+        q: asyncio.Queue[dict] = asyncio.Queue()
+        self._subscribers.append(q)
+        try:
+            await q.put({"id": self._active_id, "status": self._status})
+            while True:
+                yield await q.get()
+        finally:
+            self._subscribers.remove(q)
+
+    async def _publish(self, event: dict) -> None:
+        for q in list(self._subscribers):
+            await q.put(event)
+
+    def _free_caches(self) -> None:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        mps = getattr(torch.backends, "mps", None)
+        if mps and mps.is_available():
+            try:
+                torch.mps.empty_cache()  # type: ignore[attr-defined]
+            except Exception:
+                pass

tests/conftest.py

@@ -0,0 +1,51 @@
+"""Shared test fixtures."""
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from server.schemas import Lang, ParamSpec
+
+
+class FakeAdapter:
+    """Minimal in-memory adapter for unit tests."""
+    id = "fake"
+    label = "Fake"
+    description = "Test fake"
+    languages = [Lang(code="en", label="English")]
+    paralinguistic_tags: list[str] = ["[laugh]"]
+    supports_voice_clone = True
+    params = [ParamSpec(name="t", label="T", type="float", default=0.5, min=0.0, max=1.0)]
+
+    instances: list["FakeAdapter"] = []
+
+    def __init__(self, device: str):
+        self.device = device
+        self.loaded = False
+        self.unload_called = False
+        self.load_should_fail = False
+        FakeAdapter.instances.append(self)
+
+    def load(self) -> None:
+        if self.load_should_fail:
+            raise RuntimeError("simulated load failure")
+        self.loaded = True
+
+    def unload(self) -> None:
+        self.unload_called = True
+        self.loaded = False
+
+    def generate(self, text, reference_wav_path, language, params):
+        return (b"FAKEWAV", 24000)
+
+
+class FakeAdapterB(FakeAdapter):
+    id = "fake-b"
+    label = "Fake B"
+
+
+@pytest.fixture
+def fake_classes():
+    FakeAdapter.instances.clear()
+    return {FakeAdapter.id: FakeAdapter, FakeAdapterB.id: FakeAdapterB}

tests/test_registry.py

@@ -0,0 +1,90 @@
+import asyncio
+
+import pytest
+
+from server.registry import Registry
+
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_get_or_load_loads_first_time(fake_classes):
+    reg = Registry(adapter_classes=fake_classes, device="cpu")
+    a = await reg.get_or_load("fake")
+    assert a.loaded is True
+    assert reg.status()["status"] == "loaded"
+    assert reg.status()["id"] == "fake"
+
+
+async def test_get_or_load_reuses_active(fake_classes):
+    reg = Registry(adapter_classes=fake_classes, device="cpu")
+    a1 = await reg.get_or_load("fake")
+    a2 = await reg.get_or_load("fake")
+    assert a1 is a2
+
+
+async def test_get_or_load_swaps_to_different(fake_classes):
+    reg = Registry(adapter_classes=fake_classes, device="cpu")
+    a = await reg.get_or_load("fake")
+    b = await reg.get_or_load("fake-b")
+    assert b.loaded is True
+    assert a.unload_called is True
+    assert reg.status()["id"] == "fake-b"
+
+
+async def test_get_or_load_unknown_id_raises(fake_classes):
+    reg = Registry(adapter_classes=fake_classes, device="cpu")
+    with pytest.raises(KeyError):
+        await reg.get_or_load("nope")
+
+
+async def test_load_failure_sets_error_status(fake_classes):
+    reg = Registry(adapter_classes=fake_classes, device="cpu")
+    fake_classes["fake"].instances.clear()
+    orig_init = fake_classes["fake"].__init__
+
+    def patched_init(self, device):
+        orig_init(self, device)
+        self.load_should_fail = True
+
+    fake_classes["fake"].__init__ = patched_init
+    try:
+        with pytest.raises(RuntimeError):
+            await reg.get_or_load("fake")
+        s = reg.status()
+        assert s["status"] == "error"
+        assert "simulated" in s["last_error"]
+    finally:
+        fake_classes["fake"].__init__ = orig_init
+
+
+async def test_concurrent_activations_serialize(fake_classes):
+    reg = Registry(adapter_classes=fake_classes, device="cpu")
+    # All three resolve without error; the final state is the last requested model.
+    await asyncio.gather(
+        reg.get_or_load("fake"),
+        reg.get_or_load("fake-b"),
+        reg.get_or_load("fake"),
+    )
+    s = reg.status()
+    assert s["status"] == "loaded"
+    assert s["id"] == "fake"
+
+
+async def test_emits_loading_then_loaded_events(fake_classes):
+    reg = Registry(adapter_classes=fake_classes, device="cpu")
+    seen: list[dict] = []
+
+    async def collect():
+        async for evt in reg.stream_events():
+            seen.append(evt)
+            if evt["status"] == "loaded":
+                return
+
+    consumer = asyncio.create_task(collect())
+    await asyncio.sleep(0)
+    await reg.get_or_load("fake")
+    await asyncio.wait_for(consumer, timeout=2)
+    statuses = [e["status"] for e in seen]
+    assert "loading" in statuses
+    assert "loaded" in statuses