diffract-core 0.2.1__py3-none-any.whl

diffract/core/compute/decorator.py

@@ -0,0 +1,125 @@
+"""Kernel decorator and registration utilities."""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from typing import TYPE_CHECKING, Any
+
+from dependency_injector.wiring import Provide, inject
+
+from .exceptions import InconsistentWiring
+from .execution import KernelApplyLevel, KernelExecutionProtocol, KernelRestrictions
+from .metadata import KernelInfo
+from .registry import KernelRegistry
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+# Registration manifest for the built-in kernels. Import side effects only
+# run once per process, but each fresh registry still needs the built-ins;
+# register_default_kernels replays this manifest into the target registry.
+_DEFAULT_KERNEL_SPECS: list[dict[str, Any]] = []
+_COLLECTING_DEFAULTS = False
+
+
+@inject
+def kernel(
+    _func: Callable[..., Any] | None = None,
+    *,
+    registry: KernelRegistry = Provide["kernel_registry"],
+    name: str | None = None,
+    require_fields: tuple[str, ...] | None = None,
+    produce_fields: tuple[str, ...] | None = None,
+    apply_level: KernelApplyLevel | None = KernelApplyLevel.PARAMETER,
+    execution_protocol: KernelExecutionProtocol
+    | None = KernelExecutionProtocol.SEQUENTIAL,
+    restrictions: KernelRestrictions | None = None,
+    info: KernelInfo | None = None,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Decorator to register a function as a compute kernel.
+
+    Registers a function with the kernel registry, enabling it to be used
+    in the compute pipeline. The decorator infers required fields and default
+    configuration from the function signature unless explicitly provided.
+
+    Args:
+        _func: Function to decorate (used internally).
+        registry: Kernel registry for registration (injected).
+        name: Custom kernel name (defaults to function name).
+        require_fields: Input field names (inferred from signature if None).
+        produce_fields: Output field names (defaults to kernel name).
+        apply_level: Level at which kernel operates.
+        execution_protocol: Execution strategy for the kernel.
+        restrictions: Optional argument restrictions.
+        info: Optional documentation metadata.
+
+    Returns:
+        Decorator function or decorated function.
+
+    Raises:
+        InconsistentWiring: If registry injection is not properly configured.
+
+    Example:
+        >>> @kernel(name="custom_field", produce_fields=("result",))
+        ... def my_kernel(input_field: float, param: float = 1.0) -> float:
+        ...     return input_field * param
+    """
+    if not isinstance(registry, KernelRegistry):
+        raise InconsistentWiring
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        req_auto, cfg = registry._split_signature(func)  # noqa: SLF001
+        final_name = name or func.__name__
+        final_require = require_fields or req_auto
+        final_produce = produce_fields or (final_name,)
+
+        spec: dict[str, Any] = {
+            "name": final_name,
+            "require_fields": final_require,
+            "produce_fields": final_produce,
+            "implementation": func,
+            "apply_level": apply_level,
+            "execution_protocol": execution_protocol,
+            "restrictions": restrictions,
+            "config": cfg,
+            "info": info or KernelInfo(),
+        }
+        registry.register_kernel(**spec)
+        if _COLLECTING_DEFAULTS:
+            _DEFAULT_KERNEL_SPECS.append(spec)
+
+        return func
+
+    if _func is None:
+        return decorator
+    return decorator(_func)
+
+
+def register_default_kernels(registry: KernelRegistry | None = None) -> None:
+    """Register the built-in kernels, replaying them into ``registry``.
+
+    The first call imports the kernels subpackage, which registers every
+    ``@kernel``-decorated function into the currently wired registry and
+    records a manifest of the built-ins. Because module imports are cached
+    per process, later calls replay that manifest into ``registry`` so a
+    freshly created registry also receives the built-in kernels.
+
+    Args:
+        registry: Target registry for the manifest replay. When None, only
+            the import side effects apply (first call in the process).
+    """
+    global _COLLECTING_DEFAULTS  # noqa: PLW0603
+
+    module_name = f"{__package__}.kernels"
+    if module_name not in sys.modules:
+        _COLLECTING_DEFAULTS = True
+        try:
+            importlib.import_module(".kernels", package=__package__)
+        finally:
+            _COLLECTING_DEFAULTS = False
+
+    if registry is not None:
+        for spec in _DEFAULT_KERNEL_SPECS:
+            if not registry.has_kernel(spec["name"]):
+                registry.register_kernel(**spec)

diffract/core/compute/exceptions.py

@@ -0,0 +1,31 @@
+"""Kernel-related exception classes."""
+
+from __future__ import annotations
+
+
+class KernelError(Exception):
+    """Base class for kernel-related errors."""
+
+
+class DependencyNotFoundError(KernelError):
+    """Raised when a dependency is not found in the registry."""
+
+
+class CircularDependencyError(KernelError):
+    """Raised when a circular dependency is detected among kernels."""
+
+
+class InvalidConfigurationError(KernelError):
+    """Raised when kernel configuration is invalid."""
+
+
+class InconsistentWiringError(KernelError):
+    """Raised when dependency injection wiring is inconsistent."""
+
+
+class KernelExecutionError(KernelError):
+    """Raised when a kernel implementation raises during execution."""
+
+
+InvalidConfiguration = InvalidConfigurationError
+InconsistentWiring = InconsistentWiringError

diffract/core/compute/execution/__init__.py

@@ -0,0 +1,14 @@
+# Re-export for convenience
+from .enums import KernelApplyLevel, KernelExecutionProtocol, KernelRestrictions
+from .executor import KernelExecutor
+from .strategy import ExecutionStrategy, ParallelStrategy, SequentialStrategy
+
+__all__ = [
+    "ExecutionStrategy",
+    "KernelApplyLevel",
+    "KernelExecutionProtocol",
+    "KernelExecutor",
+    "KernelRestrictions",
+    "ParallelStrategy",
+    "SequentialStrategy",
+]

diffract/core/compute/execution/_types.py

@@ -0,0 +1,70 @@
+"""Internal data types for kernel execution.
+
+This module defines the data structures used internally by kernel runners
+for tracking execution state, required inputs, and pending work.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.params.interface import IParameterView
+
+    from .aggregation import AggregationContext
+
+
+@dataclass(frozen=True, slots=True)
+class AggregateInput:
+    """Description of an aggregate input required by an aggregation kernel.
+
+    Used to track which aggregate values need to be read from the
+    AggregateRepository when executing aggregation-level kernels.
+
+    Attributes:
+        field_name: Name of the aggregate field to read.
+        context_models: Model IDs participating in the aggregation.
+        context_params: Parameter names participating in the aggregation.
+    """
+
+    field_name: str
+    context_models: tuple[str, ...]
+    context_params: tuple[str, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class RequiredInputs:
+    """Required inputs for kernel execution, categorized by source.
+
+    Separates inputs into two categories:
+    - parameter_fields: Read directly from parameter proxies
+    - aggregate_inputs: Read from the aggregate repository
+
+    Attributes:
+        parameter_fields: Field names to read from parameters.
+        aggregate_inputs: Aggregate field descriptors for repository lookups.
+    """
+
+    parameter_fields: tuple[str, ...]
+    aggregate_inputs: tuple[AggregateInput, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class PendingGroup:
+    """A group of parameters pending aggregation kernel execution.
+
+    Represents a unit of work for aggregation kernels, containing all
+    information needed to execute the kernel for a specific group.
+
+    Attributes:
+        group_id: Unique identifier for this group.
+        context: Aggregation context with model/parameter scope.
+        view: Parameter view containing the group's parameters.
+        required_inputs: Inputs needed to execute the kernel.
+    """
+
+    group_id: str
+    context: AggregationContext
+    view: IParameterView
+    required_inputs: RequiredInputs

diffract/core/compute/execution/aggregation.py

@@ -0,0 +1,74 @@
+"""Aggregation and grouping helpers extracted from executor."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+
+from diffract.core.constants import (
+    format_contextual_field_name,
+    format_field_suffix,
+)
+from diffract.core.data.nn.params.interface import IParameterView
+
+type GroupsIterator = Iterable[tuple[str, IParameterView]]
+type ContextBuilder = Callable[[str, IParameterView], AggregationContext]
+
+
+@dataclass(frozen=True)
+class AggregationContext:
+    """Context for aggregated computations with contextual field naming.
+
+    Attributes:
+        models: Model identifiers participating in aggregation.
+        parameters: Parameter names participating in aggregation.
+    """
+
+    models: tuple[str, ...] | None = None
+    parameters: tuple[str, ...] | None = None
+
+    def to_field_suffix(self) -> str:
+        """Build a deterministic suffix from context identifiers."""
+        return format_field_suffix(models=self.models, params=self.parameters)
+
+    def create_field_name(self, field_name: str) -> str:
+        """Compose contextual field name using the context suffix."""
+        return format_contextual_field_name(
+            field_name, models=self.models, params=self.parameters
+        )
+
+
+def aggregate_parameters(
+    collection: IParameterView,
+) -> tuple[GroupsIterator, ContextBuilder]:
+    """Group parameters by model ID."""
+    model_ids = sorted({p.meta.model_id for p in collection})
+
+    def iter_groups() -> GroupsIterator:
+        for model_id in model_ids:
+            group = collection.filter_by_model_id(model_id)
+            yield model_id, group
+
+    def build_context(model_id: str, group: IParameterView) -> AggregationContext:
+        names = tuple(sorted(p.meta.name for p in group))
+        return AggregationContext(models=(model_id,), parameters=names)
+
+    return iter_groups(), build_context
+
+
+def aggregate_models(
+    collection: IParameterView,
+) -> tuple[GroupsIterator, ContextBuilder]:
+    """Group parameters by parameter name."""
+    param_names = sorted({p.meta.name for p in collection})
+
+    def iter_groups() -> GroupsIterator:
+        for param_name in param_names:
+            grp = collection.filter_by_name(param_name)
+            yield param_name, grp
+
+    def build_context(param_name: str, group: IParameterView) -> AggregationContext:
+        model_ids = tuple(sorted(p.meta.model_id for p in group))
+        return AggregationContext(models=model_ids, parameters=(param_name,))
+
+    return iter_groups(), build_context

diffract/core/compute/execution/aggregation_runner.py

@@ -0,0 +1,362 @@
+"""Aggregation-level kernel execution logic."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from diffract.core.data.nn.aggregates.metadata import AggregateMetadata
+
+from ._types import AggregateInput, PendingGroup, RequiredInputs
+from .aggregation import AggregationContext, aggregate_models, aggregate_parameters
+from .enums import KernelApplyLevel
+from .restrictions import apply_restrictions_filter
+from .strategy import create_execution_strategy
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from concurrent.futures import ProcessPoolExecutor
+
+    from diffract.core.compute.registry import KernelRegistry
+    from diffract.core.data.nn.aggregates import AggregateRepository
+    from diffract.core.data.nn.params.interface import IParameterView
+    from diffract.core.parallel import ParallelContext
+
+logger = logging.getLogger(__name__)
+
+
+class AggregationKernelRunner:
+    """Executes aggregation-level kernels (IN_MODEL and CROSS_MODEL).
+
+    Handles execution of kernels that aggregate across parameters within
+    a model or across models, with support for memory-aware batching.
+    """
+
+    def __init__(
+        self,
+        registry: KernelRegistry,
+        process_pool: ProcessPoolExecutor | None,
+        parallel: ParallelContext | None,
+        aggregate_repository: AggregateRepository | None,
+    ) -> None:
+        """Initialize the aggregation kernel runner.
+
+        Args:
+            registry: Kernel registry for metadata lookup.
+            process_pool: Optional process pool for parallel execution.
+            parallel: Optional parallel context for view operations.
+            aggregate_repository: Repository for storing aggregation results.
+        """
+        self._registry = registry
+        self._process_pool = process_pool
+        self._parallel = parallel
+        self._aggregate_repository = aggregate_repository
+
+    def run(self, kernel_name: str, parameters: IParameterView) -> None:
+        """Execute kernel with aggregation across models or parameters.
+
+        Args:
+            kernel_name: Name of the registered kernel.
+            parameters: Parameter collection to process.
+        """
+        pending_groups = self._collect_pending_groups(kernel_name, parameters)
+        if not pending_groups:
+            logger.debug(
+                "Skip execution of kernel '%s': no pending groups", kernel_name
+            )
+            return
+
+        logger.info(
+            "Executing kernel '%s' (%d groups)", kernel_name, len(pending_groups)
+        )
+
+        batches = self._batch_groups_by_budget(pending_groups, parameters)
+        for batch in batches:
+            self._prefetch_batch(batch, parameters)
+            self._execute_batch(kernel_name, batch)
+
+    def _collect_pending_groups(
+        self, kernel_name: str, parameters: IParameterView
+    ) -> list[PendingGroup]:
+        """Collect groups that need kernel execution."""
+        apply_level = self._registry.get_kernel_apply_level(kernel_name)
+
+        if apply_level == KernelApplyLevel.IN_MODEL:
+            group_iter, build_context = aggregate_parameters(parameters)
+        elif apply_level == KernelApplyLevel.CROSS_MODEL:
+            group_iter, build_context = aggregate_models(parameters)
+        else:
+            msg = f"Unsupported level for grouping: {apply_level}"
+            raise ValueError(msg)
+
+        target_fields = self._registry.get_fields_kernel_produce(kernel_name)
+        pending: list[PendingGroup] = []
+
+        for group_id, group in group_iter:
+            context = build_context(group_id, group)
+
+            if self._all_fields_computed(target_fields, context):
+                logger.debug(
+                    "Skip execution of kernel '%s' for group '%s'",
+                    kernel_name,
+                    group_id,
+                )
+                continue
+
+            required_inputs = self._get_required_inputs(kernel_name, context)
+            pending.append(
+                PendingGroup(
+                    group_id=group_id,
+                    context=context,
+                    view=group,
+                    required_inputs=required_inputs,
+                )
+            )
+
+        return pending
+
+    def _all_fields_computed(
+        self, field_names: tuple[str, ...], context: AggregationContext
+    ) -> bool:
+        """Check if all target fields are already computed for the context."""
+        if self._aggregate_repository is None:
+            return False
+
+        for field in field_names:
+            uid = AggregateMetadata.create_uid_from_context(
+                field_name=field,
+                context_models=context.models or (),
+                context_params=context.parameters or (),
+            )
+            try:
+                proxy = self._aggregate_repository.get_proxy(uid)
+            except KeyError:
+                return False
+
+            if not proxy.has_field("value"):
+                return False
+
+        return True
+
+    def _get_required_inputs(
+        self, kernel_name: str, context: AggregationContext
+    ) -> RequiredInputs:
+        """Determine required inputs for executing an aggregation kernel."""
+        parameter_fields: list[str] = []
+        aggregate_inputs: list[AggregateInput] = []
+
+        for field_name in self._registry.get_fields_kernel_require(kernel_name):
+            if not self._registry.can_produce_field(field_name):
+                parameter_fields.append(field_name)
+                continue
+
+            producer = self._registry.get_kernel_producing_field(field_name)
+            apply_level = self._registry.get_kernel_apply_level(producer)
+
+            if apply_level == KernelApplyLevel.PARAMETER:
+                parameter_fields.append(field_name)
+            else:
+                aggregate_inputs.append(
+                    AggregateInput(
+                        field_name=field_name,
+                        context_models=context.models or (),
+                        context_params=context.parameters or (),
+                    )
+                )
+
+        return RequiredInputs(
+            parameter_fields=tuple(parameter_fields),
+            aggregate_inputs=tuple(aggregate_inputs),
+        )
+
+    def _batch_groups_by_budget(
+        self, groups: list[PendingGroup], base_view: IParameterView
+    ) -> list[list[PendingGroup]]:
+        """Batch groups by memory budget using iter_chunks_by_read_budget."""
+        if not groups:
+            return []
+
+        uid_to_group: dict[str, PendingGroup] = {}
+        required_fields_by_uid: dict[str, list[str]] = {}
+
+        for group in groups:
+            for uid in group.view.list_uids():
+                uid_to_group[uid] = group
+                fields = list(group.required_inputs.parameter_fields)
+                required_fields_by_uid[uid] = fields
+
+        all_uids = list(uid_to_group.keys())
+        if not all_uids:
+            return []
+
+        combined_view = base_view[all_uids]
+        chunks = list(
+            combined_view.iter_chunks_by_read_budget(
+                required_fields_by_uid=required_fields_by_uid,
+                parallel=self._parallel,
+            )
+        )
+
+        batches: list[list[PendingGroup]] = []
+        for chunk in chunks:
+            seen_groups: set[str] = set()
+            batch: list[PendingGroup] = []
+            for uid in chunk.list_uids():
+                group = uid_to_group[uid]
+                if group.group_id not in seen_groups:
+                    seen_groups.add(group.group_id)
+                    batch.append(group)
+            if batch:
+                batches.append(batch)
+
+        return batches
+
+    def _prefetch_batch(
+        self, batch: list[PendingGroup], base_view: IParameterView
+    ) -> None:
+        """Prefetch fields for a batch of groups."""
+        # Prefetch parameter fields
+        param_fields_by_uid: dict[str, list[str]] = {}
+        for group in batch:
+            for uid in group.view.list_uids():
+                param_fields_by_uid[uid] = list(group.required_inputs.parameter_fields)
+
+        if param_fields_by_uid:
+            batch_uids = list(param_fields_by_uid.keys())
+            base_view[batch_uids].prefetch_fields(
+                fields_by_uid=param_fields_by_uid,
+                parallel=self._parallel,
+            )
+
+        # Prefetch aggregate fields
+        aggregate_uids = self._collect_aggregate_uids(batch)
+        if aggregate_uids and self._aggregate_repository is not None:
+            aggregate_view = self._aggregate_repository.create_view()
+            aggregate_view[aggregate_uids].prefetch_fields(
+                fields=["value"],
+                parallel=self._parallel,
+            )
+
+    def _collect_aggregate_uids(self, batch: list[PendingGroup]) -> list[str]:
+        """Collect all aggregate UIDs needed for a batch."""
+        aggregate_uids: list[str] = []
+        for group in batch:
+            for agg_input in group.required_inputs.aggregate_inputs:
+                uid = AggregateMetadata.create_uid_from_context(
+                    field_name=agg_input.field_name,
+                    context_models=agg_input.context_models,
+                    context_params=agg_input.context_params,
+                )
+                aggregate_uids.append(uid)
+        return aggregate_uids
+
+    def _execute_batch(self, kernel_name: str, batch: list[PendingGroup]) -> None:
+        """Execute aggregation-level kernel on a batch with streaming writes."""
+        required_args = self._registry.get_fields_kernel_require(kernel_name)
+        tasks: dict[tuple[str, AggregationContext], tuple[Any, ...]] = {}
+
+        for group in batch:
+            task_args = self._build_task_args(required_args, group)
+            tasks[(group.group_id, group.context)] = tuple(task_args)
+
+        apply_restrictions_filter(
+            kernel_name,
+            tasks,
+            self._registry.get_kernel_restrictions(kernel_name),
+        )
+        if not tasks:
+            return
+
+        if self._aggregate_repository is None:
+            msg = "AggregateRepository required for aggregation kernels"
+            raise RuntimeError(msg)
+
+        with self._aggregate_repository:
+            for (group_id, context), result in self._stream_results(kernel_name, tasks):
+                logger.debug(
+                    "Set results of kernel '%s' for group '%s'",
+                    kernel_name,
+                    group_id,
+                )
+                self._store_result(kernel_name, context, result)
+
+    def _build_task_args(
+        self, required_args: tuple[str, ...], group: PendingGroup
+    ) -> list[Any]:
+        """Build argument list for a single aggregation task."""
+        task_args: list[Any] = []
+
+        for required_arg in required_args:
+            if not self._registry.can_produce_field(required_arg):
+                # Raw field from parameters
+                values = [p.get_field(required_arg) for p in group.view]
+                task_args.append(values)
+                continue
+
+            dependency = self._registry.get_kernel_producing_field(required_arg)
+            dependency_level = self._registry.get_kernel_apply_level(dependency)
+
+            if dependency_level == KernelApplyLevel.PARAMETER:
+                # Parameter-level computed field
+                values = [p.get_field(required_arg) for p in group.view]
+                task_args.append(values)
+            else:
+                # Aggregate-level computed field
+                value = self._read_aggregate_value(required_arg, group.context)
+                task_args.append(value)
+
+        return task_args
+
+    def _read_aggregate_value(
+        self, field_name: str, context: AggregationContext
+    ) -> Any:
+        """Read an aggregate value from AggregateRepository."""
+        if self._aggregate_repository is None:
+            contextual_name = context.create_field_name(field_name)
+            msg = f"AggregateRepository not configured, cannot read '{contextual_name}'"
+            raise KeyError(msg)
+
+        uid = AggregateMetadata.create_uid_from_context(
+            field_name=field_name,
+            context_models=context.models or (),
+            context_params=context.parameters or (),
+        )
+
+        try:
+            aggregate = self._aggregate_repository.get_proxy(uid)
+        except KeyError:
+            contextual_name = context.create_field_name(field_name)
+            msg = f"Missing contextual dependency '{contextual_name}'"
+            raise KeyError(msg) from None
+
+        if not aggregate.has_field("value"):
+            contextual_name = context.create_field_name(field_name)
+            msg = f"Aggregate '{contextual_name}' exists but has no value"
+            raise KeyError(msg)
+
+        return aggregate.get_field("value")
+
+    def _store_result(
+        self, kernel_name: str, context: AggregationContext, result: Any
+    ) -> None:
+        """Store kernel result in aggregate repository."""
+        normalized = self._registry.normalize_kernel_result(kernel_name, result)
+        for field_name, value in normalized.items():
+            aggregate = self._aggregate_repository.get_or_create(
+                field_name=field_name,
+                context_models=context.models or (),
+                context_params=context.parameters or (),
+            )
+            aggregate.set_field("value", value)
+
+    def _stream_results(
+        self, kernel_name: str, tasks: dict[Any, tuple[Any, ...]]
+    ) -> Iterator[tuple[Any, Any]]:
+        """Execute kernel tasks and yield results as they become available."""
+        strategy = create_execution_strategy(
+            kernel_name,
+            self._registry,
+            self._process_pool,
+        )
+        implementation = self._registry.get_kernel_implementation(kernel_name)
+        yield from strategy.execute_tasks(kernel_name, tasks, implementation)