Source code for embedl_deploy._internal.core.backend

# Copyright (C) 2026 Embedl AB

"""Backend discovery and selection."""

import importlib
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path

from embedl_deploy._internal.core.pattern import Pattern

_INTERNAL_DIR = Path(__file__).resolve().parent.parent




@dataclass(frozen=True)
class Backend:
    """A collection of patterns for a specific hardware target."""

    #: Short identifier, matching the package directory name.
    name: str
    #: Structural rewrite patterns, applied iteratively.
    conversion_patterns: Sequence[Pattern]
    #: Fusion patterns, applied in a single pass after conversions.
    fusion_patterns: Sequence[Pattern]
    #: SmoothQuant preparation patterns.
    smooth_patterns: Sequence[Pattern]
    #: Q/DQ stub insertion patterns for quantization.
    quantized_patterns: Sequence[Pattern]



class _BackendState:
    """Module-level mutable state for backend discovery and selection."""

    #: The currently selected backend.
    backend: Backend | None = None
    #: Cached discovery result.
    backends: list[Backend] | None = None

    @classmethod
    def reset(cls) -> None:
        """Clear cached discovery results and the active backend."""
        cls.backend = None
        cls.backends = None




def discover_backends() -> list[Backend]:
    """Return all installed backends.

    Results are cached after the first call.

    :returns:
        List of discovered :class:`~embedl_deploy.backend.Backend`
        instances.
    """
    backends = _BackendState.backends
    if backends is None:
        backends = []
        for entry in sorted(_INTERNAL_DIR.iterdir()):
            if (
                not entry.is_dir()
                or entry.name.startswith("_")
                or entry.name == "core"
            ):
                continue
            module_path = f"embedl_deploy._internal.{entry.name}.backend"
            try:
                mod = importlib.import_module(module_path)
            except ModuleNotFoundError as e:
                if e.name == module_path:
                    continue
                raise
            backend = getattr(mod, "BACKEND", None)
            if isinstance(backend, Backend):
                backends.append(backend)
        _BackendState.backends = backends
    return list(backends)





def get_backend() -> Backend:
    """Return the active backend, discovering it if necessary.

    If no backend has been set via :func:`set_backend`, the installed
    backends are discovered automatically.  When exactly one is found
    it becomes the active backend.

    :returns:
        The active :class:`Backend`.
    :raises RuntimeError:
        If no backends are installed, or if multiple backends are
        installed and none has been explicitly selected.
    """
    backend = _BackendState.backend
    if backend is None:
        backends = discover_backends()
        if len(backends) == 0:
            raise RuntimeError(
                "No backends found — install at least one backend"
            )
        if len(backends) > 1:
            names = ", ".join(sorted(b.name for b in backends))
            raise RuntimeError(
                f"Multiple backends found ({names}). "
                "Call set_backend() to select one."
            )
        backend = backends[0]
        _BackendState.backend = backend
    return backend





def set_backend(name: str) -> None:
    """Select the active backend by name.

    :param name:
        The name of a discovered backend (e.g. ``"tensorrt"``).
    :raises ValueError:
        If `name` does not match any installed backend.
    """
    backend_map = {b.name: b for b in discover_backends()}
    if name not in backend_map:
        available = ", ".join(sorted(backend_map)) or "(none)"
        raise ValueError(
            f"Backend {name!r} not found. Available backends: {available}"
        )
    _BackendState.backend = backend_map[name]