diff --git a/pyproject.toml b/pyproject.toml index b5a51ea..a84c347 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,8 +62,15 @@ venv = ".venv" strict = true files = ["src", "tests"] +[[tool.mypy.overrides]] +# These ship no py.typed marker. +module = ["openvino.*", "openvino_genai.*", "huggingface_hub.*", "yt_dlp.*"] +ignore_missing_imports = true + [tool.pylint.design] max-attributes = 15 +# Keyword-only arguments do not carry the readability cost this check guards against. +max-args = 9 [tool.pylint.main] # W0621: pytest fixtures shadow their names by design. diff --git a/src/ccn_transcribe/backends/amd_backend.py b/src/ccn_transcribe/backends/amd_backend.py new file mode 100644 index 0000000..dabbd28 --- /dev/null +++ b/src/ccn_transcribe/backends/amd_backend.py @@ -0,0 +1,65 @@ +"""AMD backend -- interface only. + +Detection and routing are real: if an AMD GPU is present the registry reports +it and explains why it is stepping down. The inference path is not implemented, +because the model cached on this machine is OpenVINO IR and cannot be loaded by +ROCm, and nothing here could be tested against real hardware. + +Implementing it means transformers + a ROCm build of torch. CTranslate2 has no +ROCm support, so this cannot share the NVIDIA path. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from ccn_transcribe import errors +from ccn_transcribe.backends import probe +from ccn_transcribe.backends.base import ToolchainStatus + +if TYPE_CHECKING: + from pathlib import Path + + import numpy as np + import numpy.typing as npt + + from ccn_transcribe.backends.base import TranscribeRequest + from ccn_transcribe.transcript import TranscriptResult + +_NOT_IMPLEMENTED = "an AMD GPU was detected, but the AMD backend is not implemented in this build" +_HINT = ( + "Implement src/ccn_transcribe/backends/amd_backend.py using transformers with a " + "ROCm build of torch (pip --index-url https://download.pytorch.org/whl/rocm6.2). " + "The cached OpenVINO IR model cannot be loaded on ROCm." +) + + +class AmdBackend: + name: ClassVar[str] = "amd" + device: ClassVar[str] = "hip" + default_model: ClassVar[str] = "openai/whisper-large-v3-turbo" + + def __init__(self, model_id: str | None = None, cache_dir: Path | None = None) -> None: + self.model_id = model_id or self.default_model + self.cache_dir = cache_dir + + @classmethod + def probe_hardware(cls) -> bool: + return probe.has_amd_gpu() + + @classmethod + def probe_toolchain(cls) -> ToolchainStatus: + return ToolchainStatus(ok=False, reason=_NOT_IMPLEMENTED) + + def load(self) -> None: + raise errors.BackendNotImplementedError(_NOT_IMPLEMENTED, hint=_HINT) + + def transcribe( + self, + pcm: npt.NDArray[np.float32], # noqa: ARG002 - fixed by the Backend protocol + request: TranscribeRequest, # noqa: ARG002 - fixed by the Backend protocol + ) -> TranscriptResult: + raise errors.BackendNotImplementedError(_NOT_IMPLEMENTED, hint=_HINT) + + def close(self) -> None: + return diff --git a/src/ccn_transcribe/backends/base.py b/src/ccn_transcribe/backends/base.py new file mode 100644 index 0000000..63f9b1c --- /dev/null +++ b/src/ccn_transcribe/backends/base.py @@ -0,0 +1,57 @@ +"""The backend interface. + +Two-phase detection is the contract: ``probe_hardware`` must import nothing +heavy, and the registry only calls ``probe_toolchain`` once hardware is confirmed. +That ordering is what keeps torch off an Intel-only machine. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar, Protocol, runtime_checkable + +if TYPE_CHECKING: + from pathlib import Path + + import numpy as np + import numpy.typing as npt + + from ccn_transcribe.transcript import TranscriptResult + + +@dataclass(frozen=True, slots=True) +class ToolchainStatus: + ok: bool + reason: str = "" + + +@dataclass(frozen=True, slots=True) +class TranscribeRequest: + language: str | None = None + task: str = "transcribe" + num_beams: int = 1 + initial_prompt: str | None = None + hotwords: str | None = None + + +@runtime_checkable +class Backend(Protocol): + name: ClassVar[str] + device: ClassVar[str] + default_model: ClassVar[str] + + def __init__(self, model_id: str | None = None, cache_dir: Path | None = None) -> None: ... + + @classmethod + def probe_hardware(cls) -> bool: ... + + @classmethod + def probe_toolchain(cls) -> ToolchainStatus: ... + + def load(self) -> None: ... + + def transcribe( + self, pcm: npt.NDArray[np.float32], request: TranscribeRequest + ) -> TranscriptResult: ... + + def close(self) -> None: ... diff --git a/src/ccn_transcribe/backends/nvidia_backend.py b/src/ccn_transcribe/backends/nvidia_backend.py new file mode 100644 index 0000000..962645b --- /dev/null +++ b/src/ccn_transcribe/backends/nvidia_backend.py @@ -0,0 +1,67 @@ +"""NVIDIA backend -- interface only. + +Detection and routing are real: if an NVIDIA GPU is present the registry reports +it and explains why it is stepping down. The inference path is not implemented, +because the model cached on this machine is OpenVINO IR and cannot be loaded by +CUDA, and nothing here could be tested against real hardware. + +Implementing it means CTranslate2 (faster-whisper) with a CTranslate2 model; note +that CTranslate2 has no ROCm support, so AMD is a genuinely separate path. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from ccn_transcribe import errors +from ccn_transcribe.backends import probe +from ccn_transcribe.backends.base import ToolchainStatus + +if TYPE_CHECKING: + from pathlib import Path + + import numpy as np + import numpy.typing as npt + + from ccn_transcribe.backends.base import TranscribeRequest + from ccn_transcribe.transcript import TranscriptResult + +_NOT_IMPLEMENTED = ( + "an NVIDIA GPU was detected, but the NVIDIA backend is not implemented in this build" +) +_HINT = ( + "Implement src/ccn_transcribe/backends/nvidia_backend.py using faster-whisper " + "(CTranslate2) with a CTranslate2 model. The cached OpenVINO IR model cannot " + "be loaded on CUDA." +) + + +class NvidiaBackend: + name: ClassVar[str] = "nvidia" + device: ClassVar[str] = "cuda" + default_model: ClassVar[str] = "large-v3-turbo" + + def __init__(self, model_id: str | None = None, cache_dir: Path | None = None) -> None: + self.model_id = model_id or self.default_model + self.cache_dir = cache_dir + + @classmethod + def probe_hardware(cls) -> bool: + return probe.has_nvidia_gpu() + + @classmethod + def probe_toolchain(cls) -> ToolchainStatus: + return ToolchainStatus(ok=False, reason=_NOT_IMPLEMENTED) + + def load(self) -> None: + raise errors.BackendNotImplementedError(_NOT_IMPLEMENTED, hint=_HINT) + + def transcribe( + self, + pcm: npt.NDArray[np.float32], # noqa: ARG002 - fixed by the Backend protocol + request: TranscribeRequest, # noqa: ARG002 - fixed by the Backend protocol + ) -> TranscriptResult: + raise errors.BackendNotImplementedError(_NOT_IMPLEMENTED, hint=_HINT) + + def close(self) -> None: + return diff --git a/src/ccn_transcribe/backends/openvino_backend.py b/src/ccn_transcribe/backends/openvino_backend.py new file mode 100644 index 0000000..64124cb --- /dev/null +++ b/src/ccn_transcribe/backends/openvino_backend.py @@ -0,0 +1,229 @@ +"""OpenVINO GenAI backend -- the path this machine actually runs. + +Model resolution is offline-first so a warm cache never touches the network, and +CACHE_DIR is anchored under XDG rather than the working directory. +""" + +from __future__ import annotations + +import json +import logging +import os +import time +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar + +from ccn_transcribe import errors +from ccn_transcribe.backends import probe +from ccn_transcribe.backends.base import ToolchainStatus +from ccn_transcribe.transcript import TranscriptResult, normalize_segments + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + + from ccn_transcribe.backends.base import TranscribeRequest + +log = logging.getLogger(__name__) + +DEFAULT_MODEL = "OpenVINO/whisper-large-v3-turbo-int8-ov" +DEFAULT_REVISION = "b568445dd5dc8c695bde596f8acbb4694fd6ba64" +SAMPLE_RATE = 16_000 + + +def cache_root() -> Path: + base = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache") + return Path(base) / "ccn-transcribe" / "ov_cache" + + +def language_token(code: str | None, allowed: set[str] | None = None) -> str | None: + """Whisper wants ``<|en|>``, not ``en``.""" + if code is None: + return None + token = code if code.startswith("<|") else f"<|{code}|>" + if allowed is not None and token not in allowed: + valid = ", ".join(sorted(t.strip("<|>") for t in allowed)[:12]) + raise errors.ConfigError( + f"unknown language {code!r}", + hint=f"Valid codes include: {valid}, ...", + ) + return token + + +def supported_languages(model_dir: Path) -> set[str] | None: + config = model_dir / "generation_config.json" + if not config.exists(): + return None + try: + doc = json.loads(config.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + mapping = doc.get("lang_to_id") + return set(mapping) if mapping else None + + +def _snapshot(model_id: str, revision: str | None, *, local_only: bool) -> str: + """One typed boundary around huggingface_hub, which ships no type information.""" + # Imported lazily so --help and doctor do not pay for it. + from huggingface_hub import ( # pylint: disable=import-outside-toplevel + snapshot_download, # pyright: ignore[reportUnknownVariableType] + ) + + path: object = snapshot_download(model_id, revision=revision, local_files_only=local_only) + return str(path) + + +def resolve_model(model_id: str, revision: str | None) -> Path: + """Prefer the local cache; only reach the network when we must. + + A warm cache therefore never makes a network call, which is what lets the + pipeline run fully offline. + """ + try: + cached = _snapshot(model_id, revision, local_only=True) + except Exception: # pylint: disable=broad-exception-caught + # Any failure here just means "not in the cache"; fall through to the network. + log.info("%s is not cached; downloading (~790 MB for the default)", model_id) + else: + return Path(cached) + + try: + fetched = _snapshot(model_id, revision, local_only=False) + except Exception as exc: # pylint: disable=broad-exception-caught + raise errors.ModelUnavailableError( + f"could not obtain model {model_id}", + hint=f"Fetch it manually with: hf download {model_id}", + ) from exc + return Path(fetched) + + +class _OpenvinoBackend: + name: ClassVar[str] = "openvino" + device: ClassVar[str] = "CPU" + default_model: ClassVar[str] = DEFAULT_MODEL + + def __init__(self, model_id: str | None = None, cache_dir: Path | None = None) -> None: + self.model_id = model_id or self.default_model + self.revision = DEFAULT_REVISION if self.model_id == DEFAULT_MODEL else None + self.cache_dir = cache_dir or (cache_root() / self.device.lower()) + self.pipe: Any | None = None + self.languages: set[str] | None = None + + @classmethod + def probe_hardware(cls) -> bool: + return True + + @classmethod + def probe_toolchain(cls) -> ToolchainStatus: + try: + import openvino as ov # pylint: disable=import-outside-toplevel + except ImportError as exc: + return ToolchainStatus(ok=False, reason=f"openvino is not importable: {exc}") + devices = ov.Core().available_devices + if not cls.device_present(devices): + return ToolchainStatus( + ok=False, + reason=( + f"OpenVINO reports {devices}, without {cls.device}. The driver is " + "fine or clinfo would have failed; this is a plugin problem." + ), + ) + return ToolchainStatus(ok=True) + + @classmethod + def device_present(cls, devices: list[str]) -> bool: + return cls.device in devices + + def load(self) -> None: + # Lazy: constructing a pipeline is the only thing that needs this. + import openvino_genai as ov_genai # pylint: disable=import-outside-toplevel + + model_dir = resolve_model(self.model_id, self.revision) + self.languages = supported_languages(model_dir) + self.cache_dir.mkdir(parents=True, exist_ok=True) + try: + self.pipe = ov_genai.WhisperPipeline( + str(model_dir), self.device, CACHE_DIR=str(self.cache_dir) + ) + except Exception as exc: # pylint: disable=broad-exception-caught + raise errors.BackendRuntimeError( + f"OpenVINO could not build a pipeline on {self.device}: {exc}", + hint=self.load_hint(), + ) from exc + + def load_hint(self) -> str: + if self.device != "GPU": + return "Try --device cpu, or run `ccn-transcribe doctor`." + nodes = [n for n in probe.render_nodes() if n.vendor == "intel"] + if nodes and not nodes[0].writable: + return ( + f"{nodes[0].path} is not writable by this user. " + "Run: sudo usermod -aG render $USER, then log out and back in." + ) + return "Run `ccn-transcribe doctor` to separate driver problems from plugin ones." + + def transcribe( + self, pcm: npt.NDArray[np.float32], request: TranscribeRequest + ) -> TranscriptResult: + if self.pipe is None: + self.load() + assert self.pipe is not None # noqa: S101 - load() raises otherwise + kwargs: dict[str, Any] = { + "task": request.task, + "return_timestamps": True, + "num_beams": request.num_beams, + } + token = language_token(request.language, self.languages) + if token is not None: + kwargs["language"] = token + if request.initial_prompt: + kwargs["initial_prompt"] = request.initial_prompt + if request.hotwords: + kwargs["hotwords"] = request.hotwords + + started = time.perf_counter() + try: + raw = self.pipe.generate(pcm, **kwargs) + except Exception as exc: # pylint: disable=broad-exception-caught + raise errors.BackendRuntimeError( + f"{self.name}/{self.device} failed during generation: {exc}", + hint=self.load_hint(), + ) from exc + wall = time.perf_counter() - started + + audio_s = len(pcm) / SAMPLE_RATE + raw_chunks: list[Any] = list(getattr(raw, "chunks", None) or []) + chunks: list[tuple[float | None, float | None, str]] = [ + (c.start_ts, c.end_ts, c.text) for c in raw_chunks + ] + segments = normalize_segments(chunks, audio_s) if chunks else () + return TranscriptResult( + segments=segments, + language=request.language, + backend=self.name, + device=self.device, + model_id=self.model_id, + model_revision=self.revision, + audio_s=audio_s, + wall_s=wall, + ) + + def close(self) -> None: + self.pipe = None + + +class OpenvinoGpuBackend(_OpenvinoBackend): + device: ClassVar[str] = "GPU" + + @classmethod + def probe_hardware(cls) -> bool: + return probe.has_intel_gpu() + + @classmethod + def device_present(cls, devices: list[str]) -> bool: + # Multi-GPU systems enumerate GPU.0, GPU.1, ... + return any(d == "GPU" or d.startswith("GPU.") for d in devices) + + +class OpenvinoCpuBackend(_OpenvinoBackend): + device: ClassVar[str] = "CPU" diff --git a/src/ccn_transcribe/backends/probe.py b/src/ccn_transcribe/backends/probe.py new file mode 100644 index 0000000..53db9f5 --- /dev/null +++ b/src/ccn_transcribe/backends/probe.py @@ -0,0 +1,90 @@ +"""Hardware probing: cheap, and it imports nothing heavier than pathlib. + +The registry calls these before any toolchain import, which is what structurally +guarantees torch is never imported on a machine with no NVIDIA or AMD GPU. + +Detection reads each render node's *bound driver*, not loaded kernel modules: +/sys/module/xe can exist with zero bound devices while i915 actually owns the +card, so a module-presence check false-positives. +""" + +from __future__ import annotations + +import os +import stat +from dataclasses import dataclass +from pathlib import Path + +VENDORS = {"0x8086": "intel", "0x10de": "nvidia", "0x1002": "amd"} +INTEL_DRIVERS = ("i915", "xe") + + +@dataclass(frozen=True, slots=True) +class SysfsRoots: + drm: Path + dev_dri: Path + kfd: Path + + +DEFAULT_ROOTS = SysfsRoots( + drm=Path("/sys/class/drm"), + dev_dri=Path("/dev/dri"), + kfd=Path("/dev/kfd"), +) + +DEFAULT_PROC_NVIDIA = Path("/proc/driver/nvidia/version") + + +@dataclass(frozen=True, slots=True) +class RenderNode: + path: Path + vendor: str + driver: str + mode: int + writable: bool + + +def render_nodes(roots: SysfsRoots = DEFAULT_ROOTS) -> list[RenderNode]: + if not roots.drm.is_dir(): + return [] + nodes: list[RenderNode] = [] + for link in sorted(roots.drm.glob("renderD*")): + try: + vendor_id = (link / "device" / "vendor").read_text().strip() + driver = (link / "device" / "driver").resolve().name + except OSError: + continue + device = roots.dev_dri / link.name + mode = 0 + writable = False + if device.exists(): + mode = stat.S_IMODE(device.stat().st_mode) + writable = os.access(device, os.R_OK | os.W_OK) + nodes.append( + RenderNode( + path=device, + vendor=VENDORS.get(vendor_id, vendor_id), + driver=driver, + mode=mode, + writable=writable, + ) + ) + return nodes + + +def has_intel_gpu(roots: SysfsRoots = DEFAULT_ROOTS) -> bool: + return any(n.vendor == "intel" and n.driver in INTEL_DRIVERS for n in render_nodes(roots)) + + +def has_amd_gpu(roots: SysfsRoots = DEFAULT_ROOTS) -> bool: + if roots.kfd.exists(): + return True + return any(n.vendor == "amd" for n in render_nodes(roots)) + + +def has_nvidia_gpu( + roots: SysfsRoots = DEFAULT_ROOTS, proc_nvidia: Path = DEFAULT_PROC_NVIDIA +) -> bool: + if proc_nvidia.exists(): + return True + return any(n.vendor == "nvidia" for n in render_nodes(roots)) diff --git a/src/ccn_transcribe/backends/registry.py b/src/ccn_transcribe/backends/registry.py new file mode 100644 index 0000000..8f4a2a7 --- /dev/null +++ b/src/ccn_transcribe/backends/registry.py @@ -0,0 +1,151 @@ +"""Backend selection: detection chain, runtime fallback, and demotion.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from ccn_transcribe import errors +from ccn_transcribe.backends.amd_backend import AmdBackend +from ccn_transcribe.backends.nvidia_backend import NvidiaBackend +from ccn_transcribe.backends.openvino_backend import OpenvinoCpuBackend, OpenvinoGpuBackend + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + + import numpy as np + import numpy.typing as npt + + from ccn_transcribe.backends.base import Backend, TranscribeRequest + from ccn_transcribe.transcript import TranscriptResult + +log = logging.getLogger(__name__) + +DEFAULT_CHAIN: tuple[type[Backend], ...] = ( + NvidiaBackend, + AmdBackend, + OpenvinoGpuBackend, + OpenvinoCpuBackend, +) + +# Process-wide: once a backend fails, a 50-job batch must not retry it 50 times. +_DEMOTED: set[str] = set() + + +def key(backend: type[Backend]) -> str: + return f"{backend.name}/{backend.device}" + + +def demote(backend: type[Backend]) -> None: + _DEMOTED.add(key(backend)) + + +def reset_demotions() -> None: + _DEMOTED.clear() + + +def _select_requested( + requested: str, chain: tuple[type[Backend], ...] +) -> tuple[type[Backend], ...]: + matches = tuple(b for b in chain if requested in (b.name, b.device.lower(), key(b).lower())) + if not matches: + names = ", ".join(sorted({b.name for b in chain} | {"auto"})) + raise errors.ConfigError(f"unknown backend {requested!r}", hint=f"Choose one of: {names}") + return matches + + +def candidates( + requested: str | None = None, chain: tuple[type[Backend], ...] = DEFAULT_CHAIN +) -> Iterator[type[Backend]]: + """Backends worth trying, in order.""" + explicit = requested is not None and requested != "auto" + pool = _select_requested(requested, chain) if explicit and requested else chain + + for backend in pool: + label = key(backend) + if label in _DEMOTED: + continue + if not backend.probe_hardware(): + log.debug("skipping %s: no hardware", label) + continue + # Only now may a backend import its toolchain. + status = backend.probe_toolchain() + if not status.ok: + log.warning("skipping %s: %s", label, status.reason) + continue + yield backend + + +class BackendHolder: + """Keeps a loaded backend alive across jobs (~7s cold, ~0.7s warm).""" + + def __init__(self) -> None: + self._loaded: dict[str, Backend] = {} + + def get( + self, + backend: type[Backend], + model_id: str | None = None, + cache_dir: Path | None = None, + ) -> Backend: + label = key(backend) + instance = self._loaded.get(label) + if instance is None: + instance = backend(model_id, cache_dir) + instance.load() + self._loaded[label] = instance + return instance + + def drop(self, backend: type[Backend]) -> None: + instance = self._loaded.pop(key(backend), None) + if instance is not None: + instance.close() + + def close_all(self) -> None: + for instance in self._loaded.values(): + instance.close() + self._loaded.clear() + + +def transcribe_with_fallback( + pcm: npt.NDArray[np.float32], + request: TranscribeRequest, + *, + holder: BackendHolder | None = None, + requested: str | None = None, + model_id: str | None = None, + cache_dir: Path | None = None, + chain: tuple[type[Backend], ...] = DEFAULT_CHAIN, +) -> TranscriptResult: + """Transcribe, stepping down the chain on any runtime failure. + + Construction and the first generate() are both inside the try: a render-node + permission problem and an OpenCL JIT failure surface at one of those points, + never at device enumeration. + """ + owned = holder or BackendHolder() + explicit = requested is not None and requested != "auto" + last: Exception | None = None + + for backend in candidates(requested, chain): + label = key(backend) + try: + instance = owned.get(backend, model_id, cache_dir) + return instance.transcribe(pcm, request) + except Exception as exc: # pylint: disable=broad-exception-caught + # Any failure -- construction or first generate -- means step down. + owned.drop(backend) + if explicit: + raise errors.TranscribeError( + f"{label} failed and --device was explicit, so no fallback was tried", + hint=str(exc), + ) from exc + demote(backend) + log.warning("backend %s failed (%s); demoting it and falling back", label, exc) + last = exc + + raise errors.TranscribeError( + "no usable backend remains", + hint=str(last) if last else "Run `ccn-transcribe doctor` to see what was detected.", + ) from last diff --git a/tests/test_backends_openvino.py b/tests/test_backends_openvino.py new file mode 100644 index 0000000..850b707 --- /dev/null +++ b/tests/test_backends_openvino.py @@ -0,0 +1,298 @@ +from __future__ import annotations + +import builtins +import logging +import os +from pathlib import Path +from typing import Any + +import numpy as np +import openvino_genai +import pytest + +from ccn_transcribe import errors +from ccn_transcribe.backends import openvino_backend as ovb +from ccn_transcribe.backends import probe as probe_mod +from ccn_transcribe.backends.base import TranscribeRequest + +LIVE = pytest.mark.skipif( + not os.environ.get("CCN_LIVE"), reason="set CCN_LIVE=1 to run against the real GPU" +) + + +class TestLanguageToken: + def test_bare_code_becomes_a_token(self) -> None: + assert ovb.language_token("en") == "<|en|>" + + def test_an_existing_token_is_left_alone(self) -> None: + assert ovb.language_token("<|de|>") == "<|de|>" + + def test_none_stays_none_for_autodetect(self) -> None: + assert ovb.language_token(None) is None + + def test_an_unknown_code_is_rejected_with_suggestions(self) -> None: + with pytest.raises(errors.ConfigError, match="unknown language"): + ovb.language_token("zz", allowed={"<|en|>", "<|de|>"}) + + def test_a_known_code_passes_validation(self) -> None: + assert ovb.language_token("en", allowed={"<|en|>"}) == "<|en|>" + + +class TestSupportedLanguages: + def test_none_when_the_config_is_absent(self, tmp_path: Path) -> None: + assert ovb.supported_languages(tmp_path) is None + + def test_none_when_the_config_is_unreadable(self, tmp_path: Path) -> None: + (tmp_path / "generation_config.json").write_text("{broken") + assert ovb.supported_languages(tmp_path) is None + + def test_none_when_there_is_no_language_map(self, tmp_path: Path) -> None: + (tmp_path / "generation_config.json").write_text('{"max_length": 448}') + assert ovb.supported_languages(tmp_path) is None + + def test_reads_the_language_map(self, tmp_path: Path) -> None: + (tmp_path / "generation_config.json").write_text('{"lang_to_id": {"<|en|>": 1}}') + assert ovb.supported_languages(tmp_path) == {"<|en|>"} + + def test_reads_the_real_cached_model(self) -> None: + languages = ovb.supported_languages(ovb.resolve_model(ovb.DEFAULT_MODEL, None)) + assert languages is not None + assert "<|en|>" in languages + assert len(languages) > 90 + + +class TestCacheRoot: + def test_respects_xdg_cache_home(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + assert ovb.cache_root() == tmp_path / "ccn-transcribe/ov_cache" + + def test_falls_back_to_home_cache(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("XDG_CACHE_HOME", raising=False) + assert ovb.cache_root() == Path.home() / ".cache/ccn-transcribe/ov_cache" + + def test_is_not_relative_to_the_working_directory(self) -> None: + # The benchmark scripts use a relative .ov_cache_* and only work from the + # repo root; the pipeline must not inherit that. + assert ovb.cache_root().is_absolute() + + +class TestResolveModel: + def test_finds_the_cached_model_without_network(self) -> None: + path = ovb.resolve_model(ovb.DEFAULT_MODEL, ovb.DEFAULT_REVISION) + assert (path / "openvino_encoder_model.xml").exists() + + def test_a_missing_model_raises_with_a_manual_command( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def boom(*_a: object, **_k: object) -> None: + raise OSError("no such model") + + monkeypatch.setattr(ovb, "_snapshot", boom) + with pytest.raises(errors.ModelUnavailableError) as caught: + ovb.resolve_model("nope/nope", None) + assert "hf download" in (caught.value.hint or "") + + +class TestDeviceSelection: + def test_gpu_backend_probes_intel_hardware(self) -> None: + assert ovb.OpenvinoGpuBackend.probe_hardware() is True + + def test_cpu_backend_always_has_hardware(self) -> None: + assert ovb.OpenvinoCpuBackend.probe_hardware() is True + + def test_gpu_toolchain_is_available_here(self) -> None: + assert ovb.OpenvinoGpuBackend.probe_toolchain().ok is True + + def test_cpu_toolchain_is_available_here(self) -> None: + assert ovb.OpenvinoCpuBackend.probe_toolchain().ok is True + + def test_multi_gpu_enumeration_is_accepted(self) -> None: + assert ovb.OpenvinoGpuBackend.device_present(["CPU", "GPU.0", "GPU.1"]) is True + + def test_a_cpu_only_enumeration_is_rejected_for_gpu(self) -> None: + assert ovb.OpenvinoGpuBackend.device_present(["CPU"]) is False + + def test_a_missing_gpu_plugin_is_reported_as_a_plugin_problem( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def absent(_cls: type[ovb.OpenvinoGpuBackend], _devices: list[str]) -> bool: + return False + + monkeypatch.setattr(ovb.OpenvinoGpuBackend, "device_present", classmethod(absent)) + status = ovb.OpenvinoGpuBackend.probe_toolchain() + assert status.ok is False + assert "plugin problem" in status.reason + + def test_an_unimportable_openvino_is_reported(self, monkeypatch: pytest.MonkeyPatch) -> None: + real = builtins.__import__ + + def no_openvino(name: str, *a: object, **k: object) -> Any: + if name == "openvino": + raise ImportError("gone") + return real(name, *a, **k) # type: ignore[arg-type] + + monkeypatch.setattr(builtins, "__import__", no_openvino) + assert ovb.OpenvinoCpuBackend.probe_toolchain().ok is False + + +class TestTranscribeWiring: + class _Pipe: + def __init__(self, chunks: list[Any] | None = None, boom: bool = False) -> None: + self.chunks = chunks + self.boom = boom + self.seen: dict[str, Any] = {} + + def generate(self, _pcm: object, **kwargs: object) -> Any: + if self.boom: + raise RuntimeError("kernel compile failed") + self.seen = dict(kwargs) + return type("R", (), {"chunks": self.chunks})() + + def _chunk(self, start: float, end: float, text: str) -> Any: + return type("C", (), {"start_ts": start, "end_ts": end, "text": text})() + + def _backend(self, pipe: object) -> ovb.OpenvinoCpuBackend: + backend = ovb.OpenvinoCpuBackend() + backend.pipe = pipe + backend.languages = {"<|en|>"} + return backend + + def test_passes_timestamps_and_task(self) -> None: + pipe = self._Pipe(chunks=[]) + self._backend(pipe).transcribe(np.zeros(16000, np.float32), TranscribeRequest()) + assert pipe.seen["return_timestamps"] is True + assert pipe.seen["task"] == "transcribe" + + def test_converts_the_language_to_token_form(self) -> None: + pipe = self._Pipe(chunks=[]) + self._backend(pipe).transcribe( + np.zeros(16000, np.float32), TranscribeRequest(language="en") + ) + assert pipe.seen["language"] == "<|en|>" + + def test_omits_language_when_autodetecting(self) -> None: + pipe = self._Pipe(chunks=[]) + self._backend(pipe).transcribe(np.zeros(16000, np.float32), TranscribeRequest()) + assert "language" not in pipe.seen + + def test_forwards_prompt_and_hotwords_only_when_set(self) -> None: + pipe = self._Pipe(chunks=[]) + self._backend(pipe).transcribe( + np.zeros(16000, np.float32), + TranscribeRequest(initial_prompt="P", hotwords="H"), + ) + assert pipe.seen["initial_prompt"] == "P" + assert pipe.seen["hotwords"] == "H" + + def test_normalizes_chunks_into_segments(self) -> None: + pipe = self._Pipe(chunks=[self._chunk(0.0, -1.0, "hi")]) + got = self._backend(pipe).transcribe(np.zeros(32000, np.float32), TranscribeRequest()) + assert got.segments[0].end == pytest.approx(2.0) + + def test_reports_audio_duration_and_device(self) -> None: + pipe = self._Pipe(chunks=[self._chunk(0.0, 1.0, "hi")]) + got = self._backend(pipe).transcribe(np.zeros(32000, np.float32), TranscribeRequest()) + assert got.audio_s == pytest.approx(2.0) + assert got.device == "CPU" + assert got.backend == "openvino" + + def test_empty_chunks_give_no_segments(self) -> None: + pipe = self._Pipe(chunks=None) + assert ( + self._backend(pipe) + .transcribe(np.zeros(16000, np.float32), TranscribeRequest()) + .segments + == () + ) + + def test_a_generate_failure_becomes_a_backend_runtime_error(self) -> None: + pipe = self._Pipe(boom=True) + with pytest.raises(errors.BackendRuntimeError, match="during generation"): + self._backend(pipe).transcribe(np.zeros(16000, np.float32), TranscribeRequest()) + + def test_transcribe_loads_lazily(self, monkeypatch: pytest.MonkeyPatch) -> None: + backend = ovb.OpenvinoCpuBackend() + calls: list[int] = [] + + def fake_load(self: ovb.OpenvinoCpuBackend) -> None: + calls.append(1) + self.pipe = TestTranscribeWiring._Pipe(chunks=[]) + + monkeypatch.setattr(ovb.OpenvinoCpuBackend, "load", fake_load) + backend.transcribe(np.zeros(16000, np.float32), TranscribeRequest()) + assert calls == [1] + + def test_close_releases_the_pipeline(self) -> None: + backend = self._backend(self._Pipe()) + backend.close() + assert backend.pipe is None + + +class TestLoadHints: + def test_cpu_hint_points_at_doctor(self) -> None: + assert "doctor" in ovb.OpenvinoCpuBackend().load_hint() + + def test_gpu_hint_names_the_render_group_when_the_node_is_unwritable( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + node = probe_mod.RenderNode( + path=Path("/dev/dri/renderD128"), + vendor="intel", + driver="i915", + mode=0o660, + writable=False, + ) + + def one_node(*_a: object, **_k: object) -> list[probe_mod.RenderNode]: + return [node] + + monkeypatch.setattr(probe_mod, "render_nodes", one_node) + assert "usermod -aG render" in ovb.OpenvinoGpuBackend().load_hint() + + def test_gpu_hint_falls_back_to_doctor_when_permissions_are_fine( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def no_nodes(*_a: object, **_k: object) -> list[probe_mod.RenderNode]: + return [] + + monkeypatch.setattr(probe_mod, "render_nodes", no_nodes) + assert "doctor" in ovb.OpenvinoGpuBackend().load_hint() + + def test_a_pipeline_construction_failure_is_wrapped( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def boom(*_a: object, **_k: object) -> None: + raise RuntimeError("no device") + + monkeypatch.setattr(openvino_genai, "WhisperPipeline", boom) + with pytest.raises(errors.BackendRuntimeError, match="could not build a pipeline"): + ovb.OpenvinoCpuBackend().load() + + +@LIVE +class TestLive: + def test_real_transcription_on_the_default_device(self, sine_wav: Path) -> None: + from ccn_transcribe.media import ffmpeg + + backend = ovb.OpenvinoGpuBackend() + backend.load() + result = backend.transcribe(ffmpeg.decode_16k_mono(sine_wav), TranscribeRequest()) + assert result.device == "GPU" + + +def test_a_cache_miss_falls_through_to_the_network( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + calls: list[bool] = [] + + def snapshot(_model: str, _revision: str | None, *, local_only: bool) -> str: + calls.append(local_only) + if local_only: + raise OSError("not cached") + return str(tmp_path) + + monkeypatch.setattr(ovb, "_snapshot", snapshot) + with caplog.at_level(logging.INFO): + assert ovb.resolve_model("some/model", None) == tmp_path + assert calls == [True, False], "the cache must be consulted before the network" + assert "not cached" in caplog.text diff --git a/tests/test_backends_probe.py b/tests/test_backends_probe.py new file mode 100644 index 0000000..aff2ef7 --- /dev/null +++ b/tests/test_backends_probe.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from ccn_transcribe.backends import probe + +if TYPE_CHECKING: + from pathlib import Path + +INTEL = "0x8086" +NVIDIA = "0x10de" +AMD = "0x1002" + + +@pytest.fixture +def sysfs(tmp_path: Path) -> probe.SysfsRoots: + (tmp_path / "sys").mkdir() + (tmp_path / "dev").mkdir() + return probe.SysfsRoots(drm=tmp_path / "sys", dev_dri=tmp_path / "dev", kfd=tmp_path / "kfd") + + +def add_node(roots: probe.SysfsRoots, name: str, vendor: str, driver: str) -> None: + node = roots.drm / name + device = node / "device" + device.mkdir(parents=True) + (device / "vendor").write_text(vendor + "\n") + # The bound driver is a symlink; its basename is the driver name. + driver_dir = roots.drm.parent / "drivers" / driver + driver_dir.mkdir(parents=True, exist_ok=True) + (device / "driver").symlink_to(driver_dir) + (roots.dev_dri / name).write_text("") + + +class TestRenderNodes: + def test_empty_when_there_is_no_drm_tree(self, tmp_path: Path) -> None: + roots = probe.SysfsRoots(drm=tmp_path / "nope", dev_dri=tmp_path / "d", kfd=tmp_path / "k") + assert probe.render_nodes(roots) == [] + + def test_reads_vendor_and_bound_driver(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "i915") + [node] = probe.render_nodes(sysfs) + assert node.vendor == "intel" + assert node.driver == "i915" + + def test_unknown_vendor_id_is_passed_through(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", "0xbeef", "weird") + assert probe.render_nodes(sysfs)[0].vendor == "0xbeef" + + def test_nodes_without_a_vendor_file_are_skipped(self, sysfs: probe.SysfsRoots) -> None: + (sysfs.drm / "renderD200" / "device").mkdir(parents=True) + assert probe.render_nodes(sysfs) == [] + + def test_only_render_nodes_are_listed(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "i915") + add_node(sysfs, "card1", INTEL, "i915") + assert [n.path.name for n in probe.render_nodes(sysfs)] == ["renderD128"] + + def test_reports_permissions(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "i915") + (sysfs.dev_dri / "renderD128").chmod(0o666) + node = probe.render_nodes(sysfs)[0] + assert node.mode == 0o666 + assert node.writable is True + + def test_a_node_with_no_device_file_is_not_writable(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "i915") + (sysfs.dev_dri / "renderD128").unlink() + assert probe.render_nodes(sysfs)[0].writable is False + + +class TestVendorProbes: + def test_intel_detected_via_i915(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "i915") + assert probe.has_intel_gpu(sysfs) is True + + def test_intel_detected_via_xe(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "xe") + assert probe.has_intel_gpu(sysfs) is True + + def test_intel_not_detected_when_the_driver_is_unbound(self, sysfs: probe.SysfsRoots) -> None: + # /sys/module/xe can exist with zero bound devices, which is why the + # bound-driver symlink is the probe rather than module presence. + add_node(sysfs, "renderD128", INTEL, "vfio-pci") + assert probe.has_intel_gpu(sysfs) is False + + def test_no_intel_when_the_box_has_none(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", AMD, "amdgpu") + assert probe.has_intel_gpu(sysfs) is False + + def test_amd_detected_via_render_node(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", AMD, "amdgpu") + assert probe.has_amd_gpu(sysfs) is True + + def test_amd_detected_via_kfd(self, sysfs: probe.SysfsRoots) -> None: + sysfs.kfd.write_text("") + assert probe.has_amd_gpu(sysfs) is True + + def test_no_amd_on_an_intel_box(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "i915") + assert probe.has_amd_gpu(sysfs) is False + + def test_nvidia_detected_via_proc(self, sysfs: probe.SysfsRoots, tmp_path: Path) -> None: + proc = tmp_path / "nvidia_version" + proc.write_text("NVRM version ...") + assert probe.has_nvidia_gpu(sysfs, proc_nvidia=proc) is True + + def test_nvidia_detected_via_render_node(self, sysfs: probe.SysfsRoots) -> None: + # The proprietary stack does not guarantee a DRM render node, so sysfs is + # the fallback rather than the primary signal. + add_node(sysfs, "renderD128", NVIDIA, "nvidia-drm") + assert probe.has_nvidia_gpu(sysfs, proc_nvidia=sysfs.drm / "absent") is True + + def test_no_nvidia_on_an_intel_box(self, sysfs: probe.SysfsRoots) -> None: + add_node(sysfs, "renderD128", INTEL, "i915") + assert probe.has_nvidia_gpu(sysfs, proc_nvidia=sysfs.drm / "absent") is False + + +class TestAgainstTheRealMachine: + def test_defaults_point_at_the_real_system(self) -> None: + assert probe.DEFAULT_ROOTS.drm.name == "drm" + + def test_intel_igpu_is_detected_here(self) -> None: + # This box is an Iris Xe on i915; if this fails the probe is broken. + assert probe.has_intel_gpu() is True + + def test_no_nvidia_or_amd_here(self) -> None: + assert probe.has_nvidia_gpu() is False + assert probe.has_amd_gpu() is False diff --git a/tests/test_backends_registry.py b/tests/test_backends_registry.py new file mode 100644 index 0000000..55c413d --- /dev/null +++ b/tests/test_backends_registry.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import sys +from typing import TYPE_CHECKING, ClassVar + +import numpy as np +import pytest + +from ccn_transcribe import errors +from ccn_transcribe.backends import registry +from ccn_transcribe.backends.base import ToolchainStatus, TranscribeRequest +from ccn_transcribe.transcript import Segment, TranscriptResult + +if TYPE_CHECKING: + from pathlib import Path + + import numpy.typing as npt + +PCM = np.zeros(16_000, dtype=np.float32) +REQ = TranscribeRequest() + + +class FakeBackend: + name: ClassVar[str] = "fake" + device: ClassVar[str] = "FAKE" + default_model: ClassVar[str] = "fake-model" + + hardware: ClassVar[bool] = True + toolchain: ClassVar[ToolchainStatus] = ToolchainStatus(ok=True) + fail_on_load: ClassVar[bool] = False + fail_on_transcribe: ClassVar[bool] = False + loads: ClassVar[int] = 0 + + def __init__(self, model_id: str | None = None, cache_dir: Path | None = None) -> None: + self.model_id = model_id or self.default_model + self.cache_dir = cache_dir + self.closed = False + + @classmethod + def probe_hardware(cls) -> bool: + return cls.hardware + + @classmethod + def probe_toolchain(cls) -> ToolchainStatus: + return cls.toolchain + + def load(self) -> None: + type(self).loads += 1 + if self.fail_on_load: + raise errors.BackendRuntimeError(f"{self.name} failed to load") + + def transcribe( + self, pcm: npt.NDArray[np.float32], request: TranscribeRequest + ) -> TranscriptResult: + if self.fail_on_transcribe: + raise errors.BackendRuntimeError(f"{self.name} blew up mid-generate") + return TranscriptResult( + segments=(Segment(0.0, 1.0, self.name),), + language=request.language, + backend=self.name, + device=self.device, + model_id=self.model_id, + model_revision=None, + audio_s=len(pcm) / 16_000, + wall_s=0.1, + ) + + def close(self) -> None: + self.closed = True + + +def make_backend(label: str, **attrs: object) -> type[FakeBackend]: + return type( + f"{label}Backend", (FakeBackend,), {"name": label, "device": label.upper(), **attrs} + ) + + +@pytest.fixture(autouse=True) +def _clean() -> None: + registry.reset_demotions() + + +class TestCandidates: + def test_skips_backends_with_no_hardware(self) -> None: + a = make_backend("a", hardware=False) + b = make_backend("b") + assert list(registry.candidates(chain=(a, b))) == [b] + + def test_skips_backends_whose_toolchain_is_absent(self) -> None: + a = make_backend("a", toolchain=ToolchainStatus(ok=False, reason="not installed")) + b = make_backend("b") + assert list(registry.candidates(chain=(a, b))) == [b] + + def test_toolchain_is_never_probed_without_hardware(self) -> None: + # This ordering is what keeps torch from being imported on an Intel box. + probed: list[str] = [] + + class Guard(FakeBackend): + name: ClassVar[str] = "guard" + hardware: ClassVar[bool] = False + + @classmethod + def probe_toolchain(cls) -> ToolchainStatus: + probed.append(cls.name) + return ToolchainStatus(ok=True) + + list(registry.candidates(chain=(Guard,))) + assert probed == [] + + def test_chain_order_is_preserved(self) -> None: + a, b, c = make_backend("a"), make_backend("b"), make_backend("c") + assert list(registry.candidates(chain=(a, b, c))) == [a, b, c] + + def test_an_explicit_request_selects_only_that_backend(self) -> None: + a, b = make_backend("a"), make_backend("b") + assert list(registry.candidates(requested="b", chain=(a, b))) == [b] + + def test_an_unknown_request_is_a_config_error(self) -> None: + with pytest.raises(errors.ConfigError, match="unknown"): + list(registry.candidates(requested="nope", chain=(make_backend("a"),))) + + def test_auto_is_not_treated_as_a_name(self) -> None: + a = make_backend("a") + assert list(registry.candidates(requested="auto", chain=(a,))) == [a] + + def test_demoted_backends_are_skipped(self) -> None: + a, b = make_backend("a"), make_backend("b") + registry.demote(a) + assert list(registry.candidates(chain=(a, b))) == [b] + + +class TestFallback: + def test_uses_the_first_working_backend(self) -> None: + a, b = make_backend("a"), make_backend("b") + got = registry.transcribe_with_fallback(PCM, REQ, chain=(a, b)) + assert got.backend == "a" + + def test_falls_back_when_construction_fails(self) -> None: + a = make_backend("a", fail_on_load=True) + b = make_backend("b") + assert registry.transcribe_with_fallback(PCM, REQ, chain=(a, b)).backend == "b" + + def test_falls_back_when_the_first_generate_fails(self) -> None: + # The permission and JIT failures both surface here, not at detection. + a = make_backend("a", fail_on_transcribe=True) + b = make_backend("b") + assert registry.transcribe_with_fallback(PCM, REQ, chain=(a, b)).backend == "b" + + def test_a_failure_demotes_for_the_rest_of_the_run(self) -> None: + a = make_backend("a", fail_on_transcribe=True) + b = make_backend("b") + registry.transcribe_with_fallback(PCM, REQ, chain=(a, b)) + a.loads = 0 + registry.transcribe_with_fallback(PCM, REQ, chain=(a, b)) + assert a.loads == 0, "a 50-job batch must not retry a dead backend 50 times" + + def test_exhausting_the_chain_raises(self) -> None: + a = make_backend("a", fail_on_load=True) + with pytest.raises(errors.TranscribeError, match="no usable backend"): + registry.transcribe_with_fallback(PCM, REQ, chain=(a,)) + + def test_no_candidates_at_all_raises(self) -> None: + a = make_backend("a", hardware=False) + with pytest.raises(errors.TranscribeError, match="no usable backend"): + registry.transcribe_with_fallback(PCM, REQ, chain=(a,)) + + def test_an_explicit_device_never_falls_back(self) -> None: + # Asking for the GPU and silently getting the CPU makes a benchmark a lie. + a = make_backend("a", fail_on_transcribe=True) + b = make_backend("b") + with pytest.raises(errors.TranscribeError): + registry.transcribe_with_fallback(PCM, REQ, requested="a", chain=(a, b)) + + def test_warnings_name_the_failing_backend(self, caplog: pytest.LogCaptureFixture) -> None: + a = make_backend("a", fail_on_transcribe=True) + b = make_backend("b") + registry.transcribe_with_fallback(PCM, REQ, chain=(a, b)) + assert "a" in caplog.text.lower() + + +class TestHolder: + def test_reuses_a_loaded_backend_across_jobs(self) -> None: + a = make_backend("a") + holder = registry.BackendHolder() + holder.get(a) + holder.get(a) + assert a.loads == 1, "model load is ~7s cold; reloading per job is pure waste" + + def test_close_all_closes_instances(self) -> None: + a = make_backend("a") + holder = registry.BackendHolder() + instance = holder.get(a) + holder.close_all() + assert isinstance(instance, FakeBackend) + assert instance.closed is True + + def test_dropping_forces_a_reload(self) -> None: + a = make_backend("a") + holder = registry.BackendHolder() + holder.get(a) + holder.drop(a) + holder.get(a) + assert a.loads == 2 + + +class TestRealChain: + def test_default_chain_order(self) -> None: + assert [b.name for b in registry.DEFAULT_CHAIN] == [ + "nvidia", + "amd", + "openvino", + "openvino", + ] + + def test_cpu_is_last_and_terminal(self) -> None: + assert registry.DEFAULT_CHAIN[-1].device == "CPU" + + def test_this_machine_resolves_to_the_intel_gpu(self) -> None: + chosen = next(iter(registry.candidates())) + assert (chosen.name, chosen.device) == ("openvino", "GPU") + + def test_torch_is_never_imported_on_this_intel_only_box(self) -> None: + # The structural guarantee, made executable. + for module in ("torch", "faster_whisper", "ctranslate2", "transformers"): + sys.modules.pop(module, None) + list(registry.candidates()) + assert not {"torch", "faster_whisper", "ctranslate2"} & set(sys.modules) diff --git a/tests/test_backends_stubs.py b/tests/test_backends_stubs.py new file mode 100644 index 0000000..4f28921 --- /dev/null +++ b/tests/test_backends_stubs.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from ccn_transcribe import errors +from ccn_transcribe.backends.amd_backend import AmdBackend +from ccn_transcribe.backends.base import TranscribeRequest +from ccn_transcribe.backends.nvidia_backend import NvidiaBackend + +STUBS = [NvidiaBackend, AmdBackend] + + +@pytest.mark.parametrize("backend", STUBS) +class TestStubs: + def test_hardware_probe_is_real(self, backend: type[NvidiaBackend]) -> None: + # No NVIDIA or AMD on this machine, so detection must say so honestly. + assert backend.probe_hardware() is False + + def test_toolchain_reports_not_implemented(self, backend: type[NvidiaBackend]) -> None: + status = backend.probe_toolchain() + assert status.ok is False + assert "not implemented" in status.reason + + def test_load_raises_with_an_actionable_hint(self, backend: type[NvidiaBackend]) -> None: + with pytest.raises(errors.BackendNotImplementedError) as caught: + backend().load() + hint = caught.value.hint or "" + assert "backends/" in hint + + def test_transcribe_raises_rather_than_returning_nonsense( + self, backend: type[NvidiaBackend] + ) -> None: + with pytest.raises(errors.BackendNotImplementedError): + backend().transcribe(np.zeros(16, np.float32), TranscribeRequest()) + + def test_close_is_safe(self, backend: type[NvidiaBackend]) -> None: + backend().close() + + def test_default_model_is_not_the_openvino_ir(self, backend: type[NvidiaBackend]) -> None: + # OpenVINO IR cannot load on CUDA or ROCm, so each backend owns its default. + assert "-ov" not in backend.default_model + + def test_a_custom_model_id_is_kept(self, backend: type[NvidiaBackend]) -> None: + assert backend("custom/model").model_id == "custom/model" + + +def test_the_two_stubs_use_different_runtimes() -> None: + # CTranslate2 has no ROCm support, so these cannot share one implementation. + assert NvidiaBackend.device != AmdBackend.device