diffract-core 0.2.1__py3-none-any.whl

diffract/core/compute/execution/enums.py

@@ -0,0 +1,38 @@
+"""Kernel enumeration types and execution constants."""
+
+from __future__ import annotations
+
+from enum import Enum, Flag, auto
+
+
+class KernelApplyLevel(Enum):
+    """Defines the level at which kernel computation is applied.
+
+    PARAMETER: Apply kernel to individual parameters.
+    IN_MODEL: Apply kernel within a single model scope.
+    CROSS_MODEL: Apply kernel across multiple models.
+    """
+
+    PARAMETER = auto()
+    IN_MODEL = auto()
+    CROSS_MODEL = auto()
+
+
+class KernelExecutionProtocol(Enum):
+    """Defines the execution protocol for kernel processing.
+
+    SEQUENTIAL: Execute kernel tasks sequentially.
+    PARALLEL: Execute kernel tasks in parallel.
+    """
+
+    SEQUENTIAL = auto()
+    PARALLEL = auto()
+
+
+class KernelRestrictions(Flag):
+    """Flag-based restrictions that can be applied to kernel arguments.
+
+    BINARY: Restrict kernel to binary operation (exactly two arguments).
+    """
+
+    BINARY = auto()

diffract/core/compute/execution/executor.py

@@ -0,0 +1,125 @@
+"""Kernel execution engine and orchestration utilities.
+
+Provides the main KernelExecutor class that orchestrates kernel execution
+with dependency resolution, batching, and parallel processing support.
+"""
+
+from __future__ import annotations
+
+import logging
+from types import TracebackType
+from typing import TYPE_CHECKING, Self
+
+from .aggregation_runner import AggregationKernelRunner
+from .enums import KernelApplyLevel
+from .parameter_runner import ParameterKernelRunner
+
+if TYPE_CHECKING:
+    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 KernelExecutor:
+    """Kernel execution orchestrator with dependency resolution and batching.
+
+    Main execution engine that coordinates kernel execution across parameter
+    collections with support for dependency resolution, batching, and parallel
+    processing. Handles both individual parameter processing and aggregated
+    computations across models.
+
+    The executor delegates actual kernel execution to specialized runners:
+    - ParameterKernelRunner: handles PARAMETER-level kernels
+    - AggregationKernelRunner: handles IN_MODEL and CROSS_MODEL kernels
+
+    Memory constraints are handled via budget-based chunking in the runners.
+    """
+
+    def __init__(
+        self,
+        registry: KernelRegistry,
+        process_pool: ProcessPoolExecutor | None = None,
+        parallel: ParallelContext | None = None,
+        aggregate_repository: AggregateRepository | None = None,
+        **_kwargs: object,
+    ) -> None:
+        """Initialize kernel executor.
+
+        Args:
+            registry: Kernel registry containing registered kernels.
+            process_pool: Optional process pool for parallel kernel execution.
+            parallel: Optional parallel context for parameter view operations.
+            aggregate_repository: Optional repository for storing aggregation results.
+            **_kwargs: Additional config arguments (ignored, for DI compatibility).
+        """
+        self._registry = registry
+        self._executed: set[str] = set()
+
+        self._parameter_runner = ParameterKernelRunner(
+            registry=registry,
+            process_pool=process_pool,
+            parallel=parallel,
+        )
+        self._aggregation_runner = AggregationKernelRunner(
+            registry=registry,
+            process_pool=process_pool,
+            parallel=parallel,
+            aggregate_repository=aggregate_repository,
+        )
+
+    def __enter__(self) -> Self:
+        """Enter context manager."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit context manager, clearing execution state."""
+        self.clear_execution_state()
+
+    def clear_execution_state(self) -> None:
+        """Reset executed kernels set, allowing re-execution."""
+        self._executed.clear()
+
+    def execute(self, field_or_kernel_name: str, parameters: IParameterView) -> None:
+        """Execute a kernel or produce a field with dependency resolution.
+
+        Main entry point for kernel execution. Resolves all dependencies
+        and executes required kernels in proper order.
+
+        Args:
+            field_or_kernel_name: Kernel name or field name to execute/produce.
+            parameters: Parameter collection to process.
+        """
+        if field_or_kernel_name in self._executed:
+            return
+
+        dependencies = self._registry.resolve_dependencies(field_or_kernel_name)
+
+        for dep in dependencies:
+            if dep != field_or_kernel_name and dep not in self._executed:
+                self.execute(dep, parameters)
+
+        if self._registry.has_kernel(field_or_kernel_name):
+            self._dispatch_kernel(field_or_kernel_name, parameters)
+            self._executed.add(field_or_kernel_name)
+
+    def _dispatch_kernel(self, kernel_name: str, parameters: IParameterView) -> None:
+        """Dispatch kernel execution to the appropriate runner."""
+        apply_level = self._registry.get_kernel_apply_level(kernel_name)
+
+        if apply_level == KernelApplyLevel.PARAMETER:
+            self._parameter_runner.run(kernel_name, parameters)
+        elif apply_level in (KernelApplyLevel.IN_MODEL, KernelApplyLevel.CROSS_MODEL):
+            self._aggregation_runner.run(kernel_name, parameters)
+        else:
+            msg = f"Unknown kernel apply level: {apply_level}"
+            raise RuntimeError(msg)

diffract/core/compute/execution/parameter_runner.py

@@ -0,0 +1,125 @@
+"""Parameter-level kernel execution logic."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+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.params.interface import IParameterView
+    from diffract.core.parallel import ParallelContext
+
+logger = logging.getLogger(__name__)
+
+
+class ParameterKernelRunner:
+    """Executes parameter-level kernels with batching and streaming writes.
+
+    Handles execution of kernels that operate on individual parameters,
+    including memory-aware chunking, prefetching, and result streaming.
+    """
+
+    def __init__(
+        self,
+        registry: KernelRegistry,
+        process_pool: ProcessPoolExecutor | None,
+        parallel: ParallelContext | None,
+    ) -> None:
+        """Initialize the parameter kernel runner.
+
+        Args:
+            registry: Kernel registry for metadata lookup.
+            process_pool: Optional process pool for parallel execution.
+            parallel: Optional parallel context for view operations.
+        """
+        self._registry = registry
+        self._process_pool = process_pool
+        self._parallel = parallel
+
+    def run(self, kernel_name: str, parameters: IParameterView) -> None:
+        """Execute kernel on individual parameters.
+
+        Args:
+            kernel_name: Name of the registered kernel.
+            parameters: Parameter collection to process.
+        """
+        target_fields = self._registry.get_fields_kernel_produce(kernel_name)
+        pending = parameters.filter_by_fields(
+            *target_fields,
+            inverse_mask=True,
+            parallel=self._parallel,
+        )
+        if not pending:
+            logger.debug(
+                "Skip execution of kernel '%s': no pending parameters", kernel_name
+            )
+            return
+
+        logger.info("Executing kernel '%s'", kernel_name)
+
+        required_fields = list(self._registry.get_fields_kernel_require(kernel_name))
+        required_by_uid = {p.meta.uid: required_fields for p in pending}
+
+        chunks = list(
+            pending.iter_chunks_by_read_budget(
+                required_fields_by_uid=required_by_uid,
+                parallel=self._parallel,
+            )
+        )
+        if not chunks:
+            return
+
+        for chunk in chunks:
+            chunk.prefetch_fields(fields=required_fields, parallel=self._parallel)
+            self._execute_batch(kernel_name, chunk)
+
+    def _execute_batch(self, kernel_name: str, batch: IParameterView) -> None:
+        """Execute parameter-level kernel on a batch with streaming writes.
+
+        Results are written immediately as they become available, reducing
+        memory pressure for heavy computations (e.g., SVD matrices).
+        """
+        required_args = self._registry.get_fields_kernel_require(kernel_name)
+        tasks: dict[tuple[str, str], tuple[Any, ...]] = {}
+
+        for param in batch:
+            args = tuple(param.get_field(arg) for arg in required_args)
+            tasks[(param.meta.model_id, param.meta.uid)] = args
+
+        apply_restrictions_filter(
+            kernel_name,
+            tasks,
+            self._registry.get_kernel_restrictions(kernel_name),
+        )
+        if not tasks:
+            return
+
+        param_lookup = {(p.meta.model_id, p.meta.uid): p for p in batch}
+
+        with batch:
+            for key, result in self._stream_results(kernel_name, tasks):
+                normalized = self._registry.normalize_kernel_result(kernel_name, result)
+                param = param_lookup[key]
+                for field_name, value in normalized.items():
+                    param.set_field(field_name, 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)

diffract/core/compute/execution/restrictions.py

@@ -0,0 +1,61 @@
+"""Kernel restriction checking utilities."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Hashable
+from typing import Any
+
+from .enums import KernelRestrictions
+
+logger = logging.getLogger(__name__)
+
+_BINARY_ARG_COUNT = 2
+
+
+def apply_restrictions_filter(
+    kernel_name: str,
+    tasks: dict[Hashable, tuple[Any, ...]],
+    restrictions: KernelRestrictions | None,
+) -> None:
+    """Remove tasks that don't meet kernel restrictions (mutates tasks dict).
+
+    Args:
+        kernel_name: Name of the kernel for logging.
+        tasks: Dictionary of task_id -> args to filter in place.
+        restrictions: Optional kernel restrictions to check.
+    """
+    if not restrictions:
+        return
+
+    to_remove: list[Hashable] = []
+    for task_id, args in tasks.items():
+        if not _check_restrictions(restrictions, args):
+            logger.debug(
+                "Skip computation for kernel '%s', task '%s': restriction violated",
+                kernel_name,
+                task_id,
+            )
+            to_remove.append(task_id)
+
+    for task_id in to_remove:
+        del tasks[task_id]
+
+
+def _check_restrictions(
+    restrictions: KernelRestrictions, args: tuple[Any, ...]
+) -> bool:
+    """Check if arguments satisfy kernel restrictions.
+
+    Args:
+        restrictions: Kernel restrictions flags.
+        args: Tuple of arguments to check.
+
+    Returns:
+        True if all restrictions are satisfied, False otherwise.
+    """
+    if restrictions & KernelRestrictions.BINARY:
+        for arg in args:
+            if hasattr(arg, "__len__") and len(arg) != _BINARY_ARG_COUNT:
+                return False
+    return True

diffract/core/compute/execution/strategy.py

@@ -0,0 +1,118 @@
+"""Execution strategy implementations for kernel processing.
+
+Provides strategy pattern implementations for sequential and parallel
+kernel execution, abstracting the execution mechanism from kernel runners.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterator
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from typing import TYPE_CHECKING, Any
+
+from .enums import KernelExecutionProtocol
+from .utils import execute_kernel, marshal_kernel
+
+if TYPE_CHECKING:
+    from diffract.core.compute.registry import KernelRegistry
+
+
+class ExecutionStrategy(ABC):
+    """Abstract execution strategy interface."""
+
+    @abstractmethod
+    def execute_tasks(
+        self,
+        kernel_name: str,
+        tasks: dict[Any, tuple[Any, ...]],
+        implementation: Callable[..., Any],
+    ) -> Iterator[tuple[Any, Any]]:
+        """Execute tasks and yield (task_id, result) pairs as they complete.
+
+        Args:
+            kernel_name: Name of the kernel for error reporting.
+            tasks: Dictionary mapping task_id to argument tuples.
+            implementation: Kernel function to execute.
+
+        Yields:
+            Tuples of (task_id, result) as execution completes.
+        """
+        raise NotImplementedError
+
+
+class SequentialStrategy(ExecutionStrategy):
+    """Sequential execution strategy.
+
+    Executes tasks one by one in the current process, yielding
+    results immediately after each execution completes.
+    """
+
+    def execute_tasks(
+        self,
+        kernel_name: str,
+        tasks: dict[Any, tuple[Any, ...]],
+        implementation: Callable[..., Any],
+    ) -> Iterator[tuple[Any, Any]]:
+        """Execute tasks sequentially, yielding results one by one."""
+        for task_id, args in tasks.items():
+            result = execute_kernel(kernel_name, implementation, args)
+            yield task_id, result
+
+
+class ParallelStrategy(ExecutionStrategy):
+    """Parallel execution strategy using ProcessPoolExecutor.
+
+    Distributes tasks across worker processes and yields results
+    as they complete (not necessarily in submission order).
+    """
+
+    def __init__(self, pool: ProcessPoolExecutor) -> None:
+        """Initialize with a process pool.
+
+        Args:
+            pool: ProcessPoolExecutor for parallel execution.
+        """
+        self._pool = pool
+
+    def execute_tasks(
+        self,
+        kernel_name: str,
+        tasks: dict[Any, tuple[Any, ...]],
+        implementation: Callable[..., Any],
+    ) -> Iterator[tuple[Any, Any]]:
+        """Execute tasks in parallel, yielding results as they complete."""
+        marshaled = marshal_kernel(implementation)
+        future_to_task = {
+            self._pool.submit(execute_kernel, kernel_name, marshaled, args): task_id
+            for task_id, args in tasks.items()
+        }
+
+        for future in as_completed(future_to_task):
+            task_id = future_to_task[future]
+            yield task_id, future.result()
+
+
+def create_execution_strategy(
+    kernel_name: str,
+    registry: KernelRegistry,
+    process_pool: ProcessPoolExecutor | None,
+) -> ExecutionStrategy:
+    """Create appropriate execution strategy based on kernel config and pool.
+
+    Args:
+        kernel_name: Name of the kernel to execute.
+        registry: Kernel registry for protocol lookup.
+        process_pool: Optional process pool for parallel execution.
+
+    Returns:
+        Configured execution strategy instance.
+    """
+    if process_pool is None:
+        return SequentialStrategy()
+
+    protocol = registry.get_kernel_execution_protocol(kernel_name)
+    if protocol == KernelExecutionProtocol.PARALLEL:
+        return ParallelStrategy(process_pool)
+
+    return SequentialStrategy()

diffract/core/compute/execution/utils.py

@@ -0,0 +1,112 @@
+"""Multiprocessing utilities for kernel execution.
+
+Provides functions for marshaling, unmarshaling, and executing kernels
+in multiprocessing contexts with proper serialization handling.
+
+The marshal/unmarshal functions enable kernel functions to be transmitted
+across process boundaries for parallel execution, while execute_kernel
+provides uniform error handling across all execution modes.
+"""
+
+from __future__ import annotations
+
+import marshal
+import pickle
+import types
+from typing import TYPE_CHECKING, Any
+
+from diffract.core.compute.exceptions import KernelExecutionError
+from diffract.core.utils.exceptions import format_exception_message
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+def marshal_kernel(kernel_implementation: Callable[..., Any]) -> bytes:
+    """Serialize a Python function to bytes for cross-process transmission.
+
+    Converts a kernel function into a serialized format suitable for
+    transmission across process boundaries in multiprocessing scenarios.
+    Captures the function's code, defaults, and closure state.
+
+    Args:
+        kernel_implementation: Function to serialize.
+
+    Returns:
+        Serialized function as bytes.
+
+    Note:
+        The function must be picklable. Closures with non-serializable
+        state will fail to serialize.
+    """
+    kernel_code = kernel_implementation.__code__
+    kernel_defaults = kernel_implementation.__defaults__
+
+    kernel_closure_bytes: bytes | None = None
+    if kernel_implementation.__closure__:
+        kernel_closure = tuple(
+            c.cell_contents for c in kernel_implementation.__closure__
+        )
+        kernel_closure_bytes = pickle.dumps(kernel_closure)
+
+    return marshal.dumps((kernel_code, kernel_defaults, kernel_closure_bytes))
+
+
+def unmarshal_kernel(serialized: bytes) -> Callable[..., Any]:
+    """Reconstruct a Python function from serialized bytes.
+
+    Deserializes a kernel function from bytes back into an executable
+    function object with proper code, defaults, and closure restoration.
+
+    Args:
+        serialized: Serialized function bytes from marshal_kernel.
+
+    Returns:
+        Reconstructed function object.
+    """
+    kernel_code, kernel_defaults, kernel_closure = marshal.loads(serialized)  # noqa: S302
+    kernel_closure = (
+        pickle.loads(kernel_closure)  # noqa: S301
+        if kernel_closure
+        else None
+    )
+    if kernel_closure:
+        kernel_closure = tuple(types.CellType(obj) for obj in kernel_closure)
+    return types.FunctionType(
+        code=kernel_code,
+        globals=globals(),
+        name=None,
+        argdefs=kernel_defaults,
+        closure=kernel_closure,
+    )
+
+
+def execute_kernel(
+    kernel_name: str,
+    kernel_implementation: Callable[..., Any] | bytes,
+    kernel_args: tuple[Any, ...],
+) -> Any:
+    """Execute a kernel implementation with arguments, handling errors uniformly.
+
+    Executes a kernel function, automatically deserializing if needed,
+    and wrapping any exceptions in KernelExecutionError for consistent
+    error handling across the execution pipeline.
+
+    Args:
+        kernel_name: Name of the kernel for error reporting.
+        kernel_implementation: Function or serialized function bytes.
+        kernel_args: Arguments to pass to the kernel function.
+
+    Returns:
+        Result of kernel execution.
+
+    Raises:
+        KernelExecutionError: If kernel execution fails for any reason.
+    """
+    try:
+        if isinstance(kernel_implementation, bytes):
+            kernel_implementation = unmarshal_kernel(kernel_implementation)
+        return kernel_implementation(*kernel_args)
+    except Exception as e:
+        msg = f"Kernel '{kernel_name}' execution failed: {format_exception_message(e)}"
+        raise KernelExecutionError(msg) from e