embedl-deploy 0.1.0__tar.gz → 0.2.0__tar.gz

{embedl_deploy-0.1.0/src/embedl_deploy.egg-info → embedl_deploy-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embedl-deploy
-Version: 0.1.0
+Version: 0.2.0
 Summary: Python package to make AI models deployment-ready for any hardware.
 Author-email: Embedl AB <support@embedl.com>
 Project-URL: Homepage, https://www.embedl.com/

{embedl_deploy-0.1.0 → embedl_deploy-0.2.0}/pyproject.toml

@@ -107,5 +107,23 @@ select = [
     "PLR0402",
 ]
 
+[tool.mypy]
+ignore_missing_imports = false
+strict = true
+
+[[tool.mypy.overrides]]
+module = ["torch.*", "pytest.*"]
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = ["tests.*"]
+disallow_untyped_defs = false
+disallow_untyped_calls = false
+disable_error_code = ["misc", "no-any-return"]
+
+[[tool.mypy.overrides]]
+module = ["embedl_deploy._internal.tensorrt.modules.*"]
+disable_error_code = ["no-any-return"]
+
 [tool.setuptools.dynamic]
 version = { attr = "embedl_deploy.version.public.PUBLIC_VERSION" }

embedl_deploy-0.2.0/src/embedl_deploy/_internal/core/backend.py

@@ -0,0 +1,125 @@
+# 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."""
+
+    #: 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 quantisation.
+    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: dict[str, 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() -> dict[str, Backend]:
+    """Scan ``_internal/`` for importable backend packages.
+
+    Each subdirectory (except ``core``) is tried as
+    ``embedl_deploy._internal.<name>.backend``.  Directories whose
+    module cannot be found are skipped; import errors from
+    transitive dependencies are propagated.  Results are cached
+    after the first call.
+
+    :returns:
+        Mapping of backend name to ``Backend`` instance.
+    """
+    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[entry.name] = backend
+        _BackendState.backends = backends
+    return 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(backends))
+            raise RuntimeError(
+                f"Multiple backends found ({names}). "
+                "Call set_backend() to select one."
+            )
+        backend = next(iter(backends.values()))
+        _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.
+    """
+    backends = _discover_backends()
+    if name not in backends:
+        available = ", ".join(sorted(backends)) or "(none)"
+        raise ValueError(
+            f"Backend {name!r} not found. " f"Available backends: {available}"
+        )
+    _BackendState.backend = backends[name]

{embedl_deploy-0.1.0 → embedl_deploy-0.2.0}/src/embedl_deploy/_internal/core/match.py

@@ -9,8 +9,9 @@ occurrences of operation chains in ``torch.fx`` graphs.
 """
 
 import itertools
+import types
 
-from torch import fx, nn
+from torch import fx
 
 from embedl_deploy._internal.core.pattern import (
     Fork,
@@ -18,6 +19,7 @@ from embedl_deploy._internal.core.pattern import (
     NodeCheck,
     Pattern,
     PatternMatch,
+    SharedNodeCheck,
     Tree,
     TreeMatch,
     Trunk,
@@ -25,26 +27,40 @@ from embedl_deploy._internal.core.pattern import (
     get_module,
 )
 
-#: Module types recognized as activation functions by the matchers.
-ACTIVATION_MODULES: tuple[type[nn.Module], ...] = (
-    nn.ReLU,
-    nn.ReLU6,
-    nn.GELU,
-    nn.SiLU,
-    nn.Mish,
-    nn.Hardswish,
-    nn.Hardsigmoid,
-    nn.LeakyReLU,
-    nn.PReLU,
-    nn.ELU,
-    nn.Sigmoid,
-    nn.Tanh,
-)
+
+class _SharedNodeCheckSession:
+    """Checkpoint/rollback scope for ``SharedNodeCheck`` cache entries.
+
+    Between permutation attempts inside a fork, call :meth:`rollback`
+    explicitly.  On normal exit (success) cache entries survive so that
+    enclosing sessions can still enforce cross-fork shared-node
+    constraints.  On failure exit (no permutation matched and the block
+    falls through) ``__exit__`` rolls back automatically.
+    """
+
+    def __init__(self) -> None:
+        self._checkpoint = SharedNodeCheck.checkpoint()
+        self._succeeded = False
+
+    def __enter__(self) -> "_SharedNodeCheckSession":
+        return self
+
+    def success(self) -> None:
+        """Mark this session as successful — skip rollback on exit."""
+        self._succeeded = True
+
+    def __exit__(self, *args: object) -> None:
+        if not self._succeeded:
+            self.rollback()
+
+    def rollback(self) -> None:
+        """Reset every cache event logged since this session started."""
+        SharedNodeCheck.rollback_to(self._checkpoint)
 
 
 def _node_matches(node: fx.Node, checks: ModType | NodeCheck) -> bool:
     """Return whether `node` satisfies `checks`."""
-    if isinstance(checks, (type, tuple)):
+    if isinstance(checks, (type, types.UnionType)):
         return isinstance(get_module(node), checks)
     return checks(node)
 
@@ -130,34 +146,43 @@ def _match_fork_at(
         return None
 
     fork_node = trunk.pre_trunk_nodes[0]
-    if (
+    if getattr(fork.operator, "is_node_check", False):
+        if not fork.operator(fork_node):
+            return None
+    elif (
         fork_node.op != "call_function"
         or fork_node.target is not fork.operator
     ):
         return None
 
     args = [a for a in fork_node.args if isinstance(a, fx.Node)]
-    if len(args) != len(fork.inputs):
-        return None
-
-    for perm in itertools.permutations(range(len(fork.inputs))):
-        tree_matches: dict[int, TreeMatch] = {}
-        for arg_idx, branch_idx in enumerate(perm):
-            matched = _match_tree_at(
-                args[arg_idx],
-                fork.inputs[branch_idx],
-            )
-            if matched is None:
-                break
-            tree_matches[branch_idx] = matched
-
-        if len(tree_matches) == len(fork.inputs):
-            nested = [tree_matches[i] for i in range(len(fork.inputs))]
-            return TreeMatch(
-                pre_trunk_nodes=[fork_node],
-                trunk_nodes=trunk.trunk_nodes,
-                nested=nested,
-            )
+    perms = (
+        itertools.permutations(range(len(fork.inputs)))
+        if fork.perms_override is None
+        else fork.perms_override
+    )
+    with _SharedNodeCheckSession() as session:
+        for perm in perms:
+            session.rollback()
+            if len(perm) != len(args):
+                continue
+            fork_matched = True
+            tree_matches = [TreeMatch() for _ in fork.inputs]
+            for arg_idx, input_idx in enumerate(perm):
+                tree_match = _match_tree_at(
+                    args[arg_idx], fork.inputs[input_idx]
+                )
+                if tree_match is None:
+                    fork_matched = False
+                    break
+                tree_matches[input_idx] = tree_match
+            if fork_matched:
+                session.success()
+                return TreeMatch(
+                    pre_trunk_nodes=[fork_node],
+                    trunk_nodes=trunk.trunk_nodes,
+                    nested=tree_matches,
+                )
     return None
 
 
@@ -215,15 +240,17 @@ def match_tree(
         raise ValueError("``pattern`` has no tree to match.")
 
     matches: list[PatternMatch] = []
-    for node in reversed(graph_module.graph.nodes):
-        matched = _match_tree_at(node, tree)
-        if matched is not None:
-            matches.append(
-                PatternMatch(
-                    pattern=pattern,
-                    graph_module=graph_module,
-                    tree_match=matched,
+    with _SharedNodeCheckSession() as session:
+        for node in reversed(list(graph_module.graph.nodes)):
+            session.rollback()
+            matched = _match_tree_at(node, tree)
+            if matched is not None:
+                matches.append(
+                    PatternMatch(
+                        pattern=pattern,
+                        graph_module=graph_module,
+                        tree_match=matched,
+                    )
                 )
-            )
 
     return matches

{embedl_deploy-0.1.0 → embedl_deploy-0.2.0}/src/embedl_deploy/_internal/core/modules.py

@@ -2,11 +2,30 @@
 
 """Abstract ``nn.Module`` marker bases and tracing helpers."""
 
+# mypy: disable-error-code="misc"
+# torch lacks type stubs, so nn.Module resolves to Any.
+
 from abc import ABC
+from typing import TypeAlias
 
 from torch import fx, nn
 
-from embedl_deploy._internal.core.quantize.modules import QuantStub
+from embedl_deploy._internal.core.quantize.stubs import QuantStub
+
+ActivationLike: TypeAlias = (
+    nn.ReLU
+    | nn.ReLU6
+    | nn.GELU
+    | nn.SiLU
+    | nn.Mish
+    | nn.Hardswish
+    | nn.Hardsigmoid
+    | nn.LeakyReLU
+    | nn.PReLU
+    | nn.ELU
+    | nn.Sigmoid
+    | nn.Tanh
+)
 
 
 class ConvertedModule(nn.Module, ABC):
@@ -25,14 +44,14 @@ class FusedModule(nn.Module, ABC):
     """Marker base for all backend-specific fused modules.
 
     Backend packages (e.g. ``tensorrt``) subclass this for their concrete
-    fused modules (``FusedConvBNReLU``, etc.).  The generic Q/DQ insertion
-    pass in :mod:`~embedl_deploy._internal.core.quantize.qdq` uses
+    fused modules (``FusedConvBNAct``, etc.).  The generic Q/DQ insertion
+    pass in :mod:`~embedl_deploy._internal.core.quantize.prepare` uses
     ``isinstance(mod, FusedModule)`` to identify fused nodes without
     knowing backend-specific types.
     """
 
     #: Positional argument indices that should receive a
-    #: :class:`~embedl_deploy._internal.core.quantize.modules.QuantStub`.
+    #: :class:`~embedl_deploy._internal.core.quantize.stubs.QuantStub`.
     #: The Q/DQ insertion pass uses this to decide which inputs of the
     #: fused node to quantize.  Every subclass must set this explicitly.
     inputs_to_quantize: set[int]
@@ -40,9 +59,9 @@ class FusedModule(nn.Module, ABC):
     def __init__(self) -> None:
         super().__init__()
         #: Maps each index in :attr:`inputs_to_quantize` to a
-        #: :class:`~embedl_deploy._internal.core.quantize.modules.QuantStub`.
+        #: :class:`~embedl_deploy._internal.core.quantize.stubs.QuantStub`.
         self.input_quant_stubs: dict[int, QuantStub] = {
-            idx: QuantStub() for idx in self.inputs_to_quantize
+            idx: QuantStub({self}) for idx in self.inputs_to_quantize
         }
 
 
@@ -52,7 +71,8 @@ class _LeafTracer(fx.Tracer):
     def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool:
         if isinstance(m, (ConvertedModule, FusedModule, QuantStub)):
             return True
-        return super().is_leaf_module(m, module_qualified_name)
+        result: bool = super().is_leaf_module(m, module_qualified_name)
+        return result
 
 
 def symbolic_trace(model: nn.Module) -> fx.GraphModule: