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.patterns.main import Pattern, Phase

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


def _check_phase(
    patterns: Sequence[type[Pattern]],
    expected: Phase,
) -> None:
    for p in patterns:
        if p.phase != expected:
            raise ValueError(
                f"{p.__name__} has phase {p.phase.value!r}, "
                f"expected {expected.value!r} "
                f"(in {expected.value}_patterns)"
            )




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

    Each field holds patterns for exactly one
    :class:`~embedl_deploy._internal.core.patterns.main.Phase`.
    ``__post_init__`` validates that every pattern's phase matches the field it
    is assigned to.
    """

    #: Short unique identifier for the backend.
    name: str
    #: Structural rewrite patterns, applied iteratively.
    conversion_patterns: Sequence[type[Pattern]] = ()
    #: Fusion patterns, applied in a single pass after conversions.
    fusion_patterns: Sequence[type[Pattern]] = ()
    #: State-prep patterns, applied sequentially once each.
    state_prep_patterns: Sequence[type[Pattern]] = ()
    #: Q/DQ stub surgery patterns, looped to a fixpoint.
    stub_prep_patterns: Sequence[type[Pattern]] = ()

    def __post_init__(self) -> None:
        _check_phase(self.conversion_patterns, Phase.CONVERSION)
        _check_phase(self.fusion_patterns, Phase.FUSION)
        _check_phase(self.state_prep_patterns, Phase.STATE_PREP)
        _check_phase(self.stub_prep_patterns, Phase.STUB_PREP)



@dataclass
class _BackendState:
    """Mutable state for backend discovery and selection."""

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


_backend_state = _BackendState()




def reset_backends() -> None:
    """Clear cached discovery results and the active backend."""
    _backend_state.selected = None
    _backend_state.discovered = None





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

    Results are cached after the first call.

    :returns:
        List of discovered
        :class:`~embedl_deploy._internal.core.backend.Backend` instances.
    """
    backends = _backend_state.discovered
    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
            for attr in vars(mod).values():
                if isinstance(attr, Backend):
                    backends.append(attr)
        seen: dict[str, Backend] = {}
        for b in backends:
            if b.name in seen:
                raise ValueError(
                    "Multiple backend definitions with the same name"
                    f" {b.name!r}."
                )
            seen[b.name] = b
        _backend_state.discovered = backends
    return list(backends)





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

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

    :returns:
        The active :class:`~embedl_deploy._internal.core.backend.Backend`.
    :raises RuntimeError:
        If no backends are installed, or if multiple backends are
        installed and none has been explicitly selected.
    """
    backend = _backend_state.selected
    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]
        _backend_state.selected = 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"``,
        ``"lattice_advanced"``).
    :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}"
        )
    _backend_state.selected = backend_map[name]