diffract-core 0.2.1__py3-none-any.whl

diffract/session/field_cache.py

@@ -0,0 +1,216 @@
+"""Session-level cache for field availability.
+
+Caches the result of list_fields_by_uid() to avoid repeated expensive
+storage scans when multiple export_metrics() calls are made without
+intervening compute.apply() or other mutation operations.
+
+The cache uses a generation counter for lazy invalidation: when a mutation
+occurs, the generation increments and cached data is considered stale.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SessionFieldCache:
+    """Cache for field availability across session read operations.
+
+    This cache stores the mapping of parameter UID -> list of available fields,
+    avoiding repeated calls to list_fields_by_uid() which can be expensive
+    for large parameter collections.
+
+    The cache is automatically invalidated when mutating operations occur
+    (compute, add, erase, merge). It can also be incrementally updated
+    after compute() to avoid full rescans.
+
+    Example:
+        cache = SessionFieldCache()
+
+        # First export_metrics() populates the cache
+        fields_by_uid = view.list_fields_by_uid()
+        cache.set(fields_by_uid)
+
+        # Second export_metrics() uses cached data
+        if cache.is_valid:
+            fields_by_uid = cache.get()
+
+        # compute.apply() invalidates the cache
+        cache.invalidate()
+    """
+
+    _generation: int = field(default=0, repr=False)
+    _fields_by_uid: dict[str, list[str]] | None = field(default=None, repr=False)
+
+    def get(self) -> dict[str, list[str]] | None:
+        """Return cached field mapping or None if not populated.
+
+        Returns:
+            Dictionary mapping parameter UID to list of field names,
+            or None if cache is not populated.
+        """
+        return self._fields_by_uid
+
+    def set(self, fields_by_uid: dict[str, list[str]]) -> None:
+        """Store field mapping in cache.
+
+        Args:
+            fields_by_uid: Mapping of parameter UID to list of available fields.
+        """
+        self._fields_by_uid = fields_by_uid
+        logger.debug(
+            "Field cache populated with %d entries at generation %d",
+            len(fields_by_uid),
+            self._generation,
+        )
+
+    def update(self, fields_by_uid: dict[str, list[str]]) -> None:
+        """Merge additional field mappings into the existing cache.
+
+        Args:
+            fields_by_uid: Mapping of parameter UID to list of available
+                fields to merge into the current cache contents.
+        """
+        self._fields_by_uid.update(**fields_by_uid)
+
+    def invalidate(self) -> None:
+        """Invalidate cache (called after mutations).
+
+        Increments the generation counter and clears cached data.
+        This should be called after any operation that modifies storage:
+        compute.apply(), models.add(), models.erase(), results.erase(),
+        utils.merge_other_session().
+        """
+        self._fields_by_uid = None
+        self._generation += 1
+        logger.debug("Field cache invalidated, new generation: %d", self._generation)
+
+    def add_computed_fields(
+        self,
+        affected_uids: Iterable[str],
+        new_fields: Iterable[str],
+    ) -> None:
+        """Incrementally update cache after compute.apply() operation.
+
+        Instead of full invalidation, this method adds newly computed fields
+        to existing cache entries. This is an optimization for cases where
+        we know exactly which fields were added.
+
+        If cache is not populated, this is a no-op.
+
+        Args:
+            affected_uids: UIDs of parameters that were computed.
+            new_fields: Field names that were produced by the computation.
+        """
+        if self._fields_by_uid is None:
+            logger.debug("Skipping incremental update: cache not populated")
+            return
+
+        new_fields_set = set(new_fields)
+        updated_count = 0
+
+        for uid in affected_uids:
+            if uid in self._fields_by_uid:
+                existing = set(self._fields_by_uid[uid])
+                if not new_fields_set.issubset(existing):
+                    existing.update(new_fields_set)
+                    self._fields_by_uid[uid] = sorted(existing)
+                    updated_count += 1
+
+        logger.debug(
+            "Incrementally updated %d cache entries with fields: %s",
+            updated_count,
+            ", ".join(sorted(new_fields_set)),
+        )
+
+    def remove_fields_by_uids(
+        self, uids: Iterable[str], fields_to_remove: Iterable[str]
+    ) -> None:
+        """Remove fields from cache entries after results.erase().
+
+        Instead of full invalidation, this method removes specified fields
+        from specified cache entries. This is an optimization for results.erase().
+
+        If cache is not populated, this is a no-op.
+
+        Args:
+            uids: UIDs of cache entries to remove fields from.
+            fields_to_remove: Field names that were erased.
+        """
+        if self._fields_by_uid is None:
+            logger.debug("Skipping field removal: cache not populated")
+            return
+
+        fields_set = set(fields_to_remove)
+        updated_count = 0
+
+        for uid in uids:
+            original_len = len(self._fields_by_uid[uid])
+            self._fields_by_uid[uid] = [
+                f for f in self._fields_by_uid[uid] if f not in fields_set
+            ]
+            if len(self._fields_by_uid[uid]) != original_len:
+                updated_count += 1
+
+        logger.debug(
+            "Removed fields from %d cache entries: %s",
+            updated_count,
+            ", ".join(sorted(fields_set)),
+        )
+
+    def remove_uids(self, uids_to_remove: Iterable[str]) -> None:
+        """Remove UIDs from cache after models.erase().
+
+        Instead of full invalidation, this method removes specified UIDs
+        from the cache. This is an optimization for models.erase().
+
+        If cache is not populated, this is a no-op.
+
+        Args:
+            uids_to_remove: UIDs that were erased from storage.
+        """
+        if self._fields_by_uid is None:
+            logger.debug("Skipping UID removal: cache not populated")
+            return
+
+        removed_count = 0
+        for uid in uids_to_remove:
+            if uid in self._fields_by_uid:
+                del self._fields_by_uid[uid]
+                removed_count += 1
+
+        logger.debug("Removed %d UIDs from cache", removed_count)
+
+    @property
+    def is_valid(self) -> bool:
+        """Check if cache contains valid data.
+
+        Returns:
+            True if cache is populated and ready for use.
+        """
+        return self._fields_by_uid is not None
+
+    @property
+    def generation(self) -> int:
+        """Current generation counter.
+
+        Increments on each invalidation. Useful for debugging.
+
+        Returns:
+            Current generation number.
+        """
+        return self._generation
+
+    def clear(self) -> None:
+        """Clear cache completely and reset generation.
+
+        Use this for full session reset, not for normal invalidation.
+        """
+        self._fields_by_uid = None
+        self._generation = 0
+        logger.debug("Field cache cleared and reset")

diffract/session/namespaces/__init__.py

@@ -0,0 +1,15 @@
+"""Namespace helpers for Session facade."""
+
+from .compute import ComputeNamespace
+from .models import ModelsNamespace
+from .results import ResultsNamespace
+from .utils import UtilsNamespace
+from .viz import VizNamespace
+
+__all__ = [
+    "ComputeNamespace",
+    "ModelsNamespace",
+    "ResultsNamespace",
+    "UtilsNamespace",
+    "VizNamespace",
+]

diffract/session/namespaces/compute/__init__.py

@@ -0,0 +1,219 @@
+"""Compute namespace for Session."""
+
+from __future__ import annotations
+
+import difflib
+import logging
+from typing import TYPE_CHECKING, Any
+
+from dependency_injector.wiring import Provide, inject
+
+from diffract.core.compute import KernelConfig, KernelExecutor, KernelRegistry
+from diffract.core.compute.execution import (
+    KernelApplyLevel,
+    KernelExecutionProtocol,
+    KernelRestrictions,
+)
+from diffract.core.data.nn.params.interface import IParameterRepository
+from diffract.core.utils.exceptions import format_exception_message
+from diffract.session.errors import KernelNotFoundError, SessionError
+from diffract.session.session import Session, SessionContext
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from diffract.core.data.nn.params.schema import ParameterType
+
+logger = logging.getLogger(__name__)
+
+
+def _did_you_mean(name: str, candidates: list[str]) -> str:
+    close = difflib.get_close_matches(name, sorted(set(candidates)), n=3, cutoff=0.6)
+    return f" Did you mean: {', '.join(close)}?" if close else ""
+
+
+class ComputeNamespace:
+    """Kernel compute API for Session."""
+
+    @inject
+    def __init__(
+        self,
+        session_or_context: Session | SessionContext,
+        parameter_repository: IParameterRepository = Provide["nn.parameter_repository"],
+        kernel_registry: KernelRegistry = Provide["compute.kernel_registry"],
+        kernel_executor_factory: Callable[[], KernelExecutor] = Provide[
+            "compute.kernel_executor.provider"
+        ],
+    ) -> None:
+        self.__session_or_context = session_or_context
+        self.__param_repo = parameter_repository
+        self.__kernel_registry = kernel_registry
+        self.__kernel_executor_factory = kernel_executor_factory
+
+    def apply(
+        self,
+        *fields_to_produce: str,
+    ) -> None:
+        """Apply computational kernels to stored parameters.
+
+        Executes the specified kernels on filtered parameters in dependency order.
+        Results are automatically stored and can be retrieved using export_metrics().
+
+        Args:
+            *fields_to_produce: Names of fields to compute using registered kernels.
+
+        Raises:
+            KernelNotFoundError: If any specified field cannot be produced.
+            SessionError: If kernel execution fails.
+        """
+        with self.__session_or_context:
+            if not fields_to_produce:
+                logger.warning("No fields specified for computation")
+                return
+
+            for field_name in fields_to_produce:
+                if not self.__kernel_registry.can_produce_field(field_name):
+                    available = self.__kernel_registry.list_fields_can_produce()
+                    hint = _did_you_mean(field_name, available)
+                    msg = (
+                        f"Cannot produce '{field_name}': no registered kernel "
+                        f"produces it.{hint} Use "
+                        "session.compute.list_available_metrics() to see all fields."
+                    )
+                    raise KernelNotFoundError(msg)
+
+            if isinstance(self.__session_or_context, Session):
+                pending = self.__param_repo.create_view()
+            else:
+                pending = self.__session_or_context._param_filter_context
+
+            try:
+                with self.__kernel_executor_factory() as executor:
+                    for field_name in fields_to_produce:
+                        executor.execute(
+                            field_or_kernel_name=field_name, parameters=pending
+                        )
+
+                logger.info(
+                    "Successfully produced fields: %s", ", ".join(fields_to_produce)
+                )
+
+                self.__session_or_context._field_cache.invalidate()
+
+            except SessionError:
+                raise
+            except Exception as e:
+                msg = f"Failed to produce fields: {format_exception_message(e)}"
+                raise SessionError(msg) from e
+
+    def list_available_kernels(self, verbose: bool = False) -> list[str]:
+        """List all available computational kernels.
+
+        Args:
+            verbose: If True, include detailed kernel information including
+                dependencies and field requirements.
+
+        Returns:
+            List of kernel names or detailed kernel descriptions.
+        """
+        with self.__session_or_context:
+            return self.__kernel_registry.list_kernels(verbose=verbose)
+
+    def list_available_metrics(self, verbose: bool = False) -> list[str]:
+        """List all fields that can be computed by registered kernels.
+
+        Args:
+            verbose: If True, include detailed field information including
+                producing kernels and dependencies.
+
+        Returns:
+            List of computable field names or detailed field descriptions.
+        """
+        with self.__session_or_context:
+            return self.__kernel_registry.list_fields_can_produce(verbose=verbose)
+
+    def configure_kernel(self, kernel_name: str, **config: Any) -> None:
+        """Configure parameters for a specific kernel.
+
+        Updates the configuration for a registered kernel, affecting
+        subsequent computations using that kernel.
+
+        Args:
+            kernel_name: Name of kernel to configure.
+            **config: Kernel configuration parameters as keyword arguments.
+
+        Raises:
+            KernelNotFoundError: If kernel is not registered in the registry.
+        """
+        with self.__session_or_context:
+            if not self.__kernel_registry.has_kernel(kernel_name):
+                hint = _did_you_mean(kernel_name, self.__kernel_registry.list_kernels())
+                msg = f"Kernel '{kernel_name}' not found in registry.{hint}"
+                raise KernelNotFoundError(msg)
+
+            kernel_config = KernelConfig(**config)
+            self.__kernel_registry.configure_kernel(kernel_name, kernel_config)
+
+            logger.info("Configured kernel '%s' with: %s", kernel_name, config)
+
+    def kernel(
+        self,
+        _func: Callable[..., Any] | None = None,
+        *,
+        name: str | None = None,
+        require_fields: tuple[str, ...] | None = None,
+        produce_fields: tuple[str, ...] | None = None,
+        apply_level: KernelApplyLevel = KernelApplyLevel.PARAMETER,
+        execution_protocol: KernelExecutionProtocol = (
+            KernelExecutionProtocol.SEQUENTIAL
+        ),
+        restrictions: KernelRestrictions | None = None,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Decorator for registering custom kernels in the session registry.
+
+        Creates a kernel that can compute metrics on neural network parameters.
+        Dependencies are automatically inferred from function arguments.
+
+        Args:
+            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: Granularity at which the kernel is applied.
+            execution_protocol: Execution protocol for the kernel.
+            restrictions: Optional restrictions applied to kernel execution.
+
+        Returns:
+            Decorator function for registering the kernel.
+
+        Example:
+            >>> session = Session()
+            >>> with session:
+            ...
+            ...     @session.compute.kernel()
+            ...     def my_metric(frob_norm: float, *, scale: float = 1.0) -> float:
+            ...         return frob_norm * scale
+        """
+
+        def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+            req_auto, cfg = self.__kernel_registry._split_signature(func)
+            final_name = name or func.__name__
+            final_require = require_fields or req_auto
+            final_produce = produce_fields or (final_name,)
+
+            self.__kernel_registry.register_kernel(
+                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=None,
+            )
+            return func
+
+        if _func is None:
+            return decorator
+
+        return decorator(_func)