diffract-core 0.2.1__py3-none-any.whl

diffract/session/namespaces/models/__init__.py

@@ -0,0 +1,282 @@
+"""Model management namespace for Session."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from dataclasses import replace
+from typing import TYPE_CHECKING, Any
+
+from dependency_injector.wiring import Provide, inject
+
+from diffract.core.constants import TABLE_AGGREGATES, TABLE_PARAMETERS
+from diffract.core.data.nn.extractors.base import ExtractorOverrides, ParameterOverrides
+from diffract.core.utils.exceptions import format_exception_message
+from diffract.session.errors import (
+    ModelAlreadyExistsError,
+    ModelNotFoundError,
+    SessionError,
+)
+from diffract.session.session import Session, SessionContext
+
+from .parameters import ParametersNamespace
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from diffract.core.data.nn.aggregates.repository import AggregateRepository
+    from diffract.core.data.nn.extractors.handlers.base import ParameterHandler
+    from diffract.core.data.nn.extractors.interface import IParameterExtractor
+    from diffract.core.data.nn.params.interface import IParameterRepository
+
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "ModelsNamespace",
+    "ParametersNamespace",
+]
+
+
+class ModelsNamespace:
+    """Model management API."""
+
+    @inject
+    def __init__(
+        self,
+        session_or_context: Session | SessionContext,
+        parameter_repository: IParameterRepository = Provide["nn.parameter_repository"],
+        aggregate_repository: AggregateRepository = Provide["nn.aggregate_repository"],
+        extractor_factory: Callable[..., IParameterExtractor] = Provide[
+            "nn.extractor_factory.provider"
+        ],
+    ) -> None:
+        self.__session_or_context = session_or_context
+        self.__param_repo = parameter_repository
+        self.__agg_repo = aggregate_repository
+        self.__extractor_factory = extractor_factory
+
+        self.parameters = ParametersNamespace(session_or_context)
+
+    def add(
+        self,
+        model: Any,
+        model_id: str | None = None,
+        parameter_overrides: dict[str, ParameterOverrides] | None = None,
+        custom_handlers: Iterable[ParameterHandler] | None = None,
+    ) -> None:
+        """Add a neural network model to the session.
+
+        Extracts parameters from the provided model and stores them for analysis.
+        Parameters are automatically persisted and can be retrieved in future sessions.
+
+        Args:
+            model: Neural network model object (PyTorch, TensorFlow, etc.)
+                or a dict mapping parameter names to NumPy arrays.
+            model_id: Unique identifier for the model. Auto-generated if None.
+            parameter_overrides: Custom parameter extraction overrides by name.
+            custom_handlers: Additional parameter handlers for custom parameter types.
+
+        Raises:
+            ModelAlreadyExistsError: If model_id already exists in the session.
+            SessionError: If parameter extraction fails or yields no parameters.
+        """
+        with self.__session_or_context:
+            try:
+                extractor = self.__extractor_factory(
+                    model=model,
+                    overrides=self.__build_extractor_overrides(
+                        model_id, parameter_overrides
+                    ),
+                    custom_handlers=custom_handlers,
+                )
+
+                previous_len = len(self.__param_repo)
+                extractor.extract_parameters(self.__param_repo)
+                new_len = len(self.__param_repo)
+
+                self.__session_or_context._field_cache.invalidate()
+
+            except SessionError:
+                raise
+            except Exception as e:
+                msg = f"Failed to add model: {format_exception_message(e)}"
+                raise SessionError(msg) from e
+
+            if new_len == previous_len:
+                msg = (
+                    f"Model '{model_id or type(model).__name__}' produced no "
+                    "parameters. Unsupported parameter types are skipped during "
+                    "extraction; supported inputs are PyTorch, TensorFlow, "
+                    "Flax, and ONNX models (install the matching extra) and "
+                    "dicts of 2D floating-point NumPy arrays."
+                )
+                raise SessionError(msg)
+
+            logger.info(
+                "Added model '%s' with %d parameters",
+                model_id or "<auto-generated>",
+                new_len - previous_len,
+            )
+
+    def rename(
+        self,
+        model_id: str,
+        new_model_id: str,
+    ) -> None:
+        """Rename a model in the session.
+
+        Args:
+            model_id: ID of the model to rename.
+            new_model_id: New ID for the model.
+
+        Raises:
+            ModelNotFoundError: If model_id does not exist.
+            ModelAlreadyExistsError: If new_model_id already exists.
+        """
+        with self.__session_or_context:
+            params = self.__param_repo.create_view()
+            model_params = params.filter_by_model_id(model_id)
+
+            if not model_params:
+                msg = f"Model '{model_id}' not found"
+                raise ModelNotFoundError(msg)
+
+            new_model_params = params.filter_by_model_id(new_model_id)
+            if new_model_params:
+                msg = f"Model '{new_model_id}' already exists"
+                raise ModelAlreadyExistsError(msg)
+
+            param_uids = model_params.list_uids()
+
+            aggs = self.__agg_repo.create_view()
+            model_aggs = aggs.filter_by_context_models(model_id)
+            agg_uids = model_aggs.list_uids()
+
+            with self.__param_repo:
+                for uid in param_uids:
+                    param = params[uid]
+                    old_meta = param.meta
+                    new_meta = replace(old_meta, model_id=new_model_id)
+
+                    object.__setattr__(param, "meta", new_meta)
+
+                    meta_dict = new_meta.to_dict()
+                    meta_dict.pop("uid")
+                    self.__param_repo.metadata_index.upsert(
+                        TABLE_PARAMETERS,
+                        uid=uid,
+                        **meta_dict,
+                    )
+
+            with self.__agg_repo:
+                for uid in agg_uids:
+                    agg = aggs[uid]
+                    old_meta = agg.meta
+                    new_context_models = tuple(
+                        new_model_id if cm == model_id else cm
+                        for cm in old_meta.context_models
+                    )
+                    new_meta = replace(old_meta, context_models=new_context_models)
+
+                    object.__setattr__(agg, "meta", new_meta)
+
+                    meta_dict = new_meta.to_dict()
+                    meta_dict.pop("uid")
+                    self.__agg_repo.metadata_index.upsert(
+                        TABLE_AGGREGATES,
+                        uid=uid,
+                        **meta_dict,
+                    )
+
+            self.__session_or_context._field_cache.invalidate()
+
+            logger.info("Renamed model '%s' to '%s'", model_id, new_model_id)
+
+    def erase(
+        self,
+        *model_ids: str,
+        erase_all: bool = False,
+    ) -> None:
+        """Erase all parameters and results for specified models.
+
+        Removes model parameters from both memory and persistent storage.
+        Use with caution as this operation cannot be undone.
+
+        Args:
+            *model_ids: Model IDs to erase. Required if erase_all is False.
+            erase_all: If True, erases all models. model_ids must be empty.
+
+        Raises:
+            ValueError: If neither model_ids nor erase_all is specified,
+                or if both are specified.
+            ModelNotFoundError: If any specified model ID is not found.
+        """
+        with self.__session_or_context:
+            if (not model_ids) ^ erase_all:
+                if not model_ids:
+                    msg = "No model_ids provided and erase_all=False"
+                else:
+                    msg = "Cannot specify both model_ids and erase_all=True"
+                raise ValueError(msg)
+
+            if erase_all:
+                model_ids = tuple(self.list())
+
+            params = self.__param_repo.create_view()
+            aggs = self.__agg_repo.create_view()
+
+            erased_uids: list[str] = []
+            for model_id in model_ids:
+                model_params = params.filter_by_model_id(model_id)
+
+                if not model_params:
+                    msg = f"Model '{model_id}' not found"
+                    raise ModelNotFoundError(msg)
+
+                model_aggs = aggs.filter_by_context_models(model_id)
+
+                erased_uids.extend(model_params.list_uids())
+                model_params.clear(erase=True)
+                if model_aggs:
+                    model_aggs.clear(erase=True)
+
+                logger.info("Erased model '%s'", model_id)
+
+            self.__session_or_context._field_cache.remove_uids(erased_uids)
+
+    def list(self) -> list[str]:
+        """List all model IDs currently in the session.
+
+        Returns:
+            Sorted list of unique model IDs present in the session.
+        """
+        with self.__session_or_context:
+            if isinstance(self.__session_or_context, Session):
+                params = self.__param_repo.create_view()
+            else:
+                params = self.__session_or_context._param_filter_context
+            model_ids = {param.meta.model_id for param in params}
+            return sorted(model_ids)
+
+    def __build_extractor_overrides(
+        self,
+        model_id: str | None,
+        parameter_overrides: dict[str, ParameterOverrides] | None,
+    ) -> ExtractorOverrides:
+        extractor_overrides_kwargs: dict[str, Any] = {}
+
+        if model_id:
+            params = self.__param_repo.create_view()
+            model_params = params.filter_by_model_id(model_id)
+
+            if model_params:
+                msg = f"Model with uid {model_id} already exists"
+                raise ModelAlreadyExistsError(msg)
+
+            extractor_overrides_kwargs["model_id"] = model_id
+
+        if parameter_overrides:
+            extractor_overrides_kwargs["parameter_overrides"] = parameter_overrides
+
+        return ExtractorOverrides(**extractor_overrides_kwargs)

diffract/session/namespaces/models/parameters/__init__.py

@@ -0,0 +1,133 @@
+"""Parameter management namespace for Session models."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from dependency_injector.wiring import Provide, inject
+
+from diffract.core.parallel import ParallelContext
+from diffract.session.errors import SessionError
+from diffract.session.session import Session, SessionContext
+
+from .meta_patcher import MetadataPatcher, MetadataPatchError
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.params.interface import IParameterRepository
+
+logger = logging.getLogger(__name__)
+
+
+class ParametersNamespace:
+    """Parameter management API for Session models."""
+
+    @inject
+    def __init__(
+        self,
+        session_or_context: Session | SessionContext,
+        parameter_repository: IParameterRepository = Provide["nn.parameter_repository"],
+        parallel_context_factory: Callable[[], ParallelContext] = Provide[
+            "parallel_singleton.thread_pool_context.provider"
+        ],
+    ) -> None:
+        self.__session_or_context = session_or_context
+        self.__param_repo = parameter_repository
+        self.__parallel_context_factory = parallel_context_factory
+
+    def list(
+        self,
+        *,
+        verbose: bool = False,
+    ) -> list[dict[str, Any]]:
+        """List parameter metadata for inspection.
+
+        Provides detailed information about parameters currently stored
+        in the session, including computed fields and metadata.
+
+        Returns:
+            List of dictionaries containing parameter metadata including
+            UID, name, model ID, type, available fields, and custom metadata.
+        """
+        with self.__session_or_context:
+            if isinstance(self.__session_or_context, Session):
+                params = self.__param_repo.create_view()
+            else:
+                params = self.__session_or_context._param_filter_context
+
+            if not verbose:
+                return [
+                    {
+                        "uid": param.meta.uid,
+                        "name": param.meta.name,
+                        "model_id": param.meta.model_id,
+                        "parameter_type": param.meta.ptype.name,
+                    }
+                    for param in params
+                ]
+
+            if self.__session_or_context._field_cache.is_valid:
+                field_cache = self.__session_or_context._field_cache.get()
+            else:
+                field_cache = params.list_fields_by_uid(
+                    parallel=self.__parallel_context_factory()
+                )
+                self.__session_or_context._field_cache.set(field_cache)
+
+            return [
+                {
+                    "uid": param.meta.uid,
+                    "name": param.meta.name,
+                    "model_id": param.meta.model_id,
+                    "parameter_type": param.meta.ptype.name,
+                    "available_fields": field_cache.get(param.meta.uid, []),
+                    "other_meta": param.meta.other_meta,
+                }
+                for param in params
+            ]
+
+    def patch_meta(
+        self,
+        *,
+        updates: dict[str, Any],
+        force: bool = False,
+    ) -> None:
+        """Update custom metadata (`other_meta`) for selected parameters.
+
+        This method updates only the `other_meta` portion of parameter metadata.
+        Changes are persisted by rewriting the stored `__metadata__` field, so
+        subsequent sessions will see the updated metadata.
+
+        Args:
+            updates: Mapping of parameter UID to metadata updates. Each value
+                must be a dict of {key: value} pairs to apply to that
+                parameter's other_meta.
+                Example: {"uid1": {"dataset": "imagenet", "epoch": 10}}
+            force: If False (default), raise if any key already exists in
+                `other_meta`. If True, overwrite existing keys.
+
+        Raises:
+            SessionError: If conflicts are detected (force=False), invalid keys
+                are provided, or updates structure is invalid.
+        """
+        with self.__session_or_context:
+            if not updates:
+                return
+
+            for uid, value in updates.items():
+                if not isinstance(value, dict):
+                    msg = (
+                        f"Invalid update for uid '{uid}': expected dict[str, Any], "
+                        f"got {type(value).__name__}"
+                    )
+                    raise SessionError(msg)
+
+            param_view = self.__param_repo.create_view()
+            pending = param_view[updates.keys()]
+            patcher = MetadataPatcher()
+
+            try:
+                patcher.patch(updates=updates, parameters=pending, force=force)
+            except MetadataPatchError as e:
+                raise SessionError(str(e)) from e

diffract/session/namespaces/models/parameters/meta_patcher.py

@@ -0,0 +1,167 @@
+"""Parameter management namespace for Session models."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+from typing import TYPE_CHECKING, Any
+
+from tqdm.auto import tqdm
+
+from diffract.core.constants import (
+    PROGRESS_BAR_DELAY_SEC,
+    PROGRESS_BAR_MIN_ITEMS,
+    TABLE_PARAMETERS,
+)
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.params.interface import IParameterView
+    from diffract.core.data.nn.params.proxy import ParameterDataProxy
+
+logger = logging.getLogger(__name__)
+
+_ERROR_PREVIEW_LIMIT = 10
+
+
+class MetadataPatchError(Exception):
+    """Error during metadata patching."""
+
+
+class MetadataPatcher:
+    """Utility for patching parameter metadata.
+
+    Handles conflict detection and applies updates to the other_meta
+    portion of parameter metadata.
+
+    Example:
+        >>> patcher = MetadataPatcher()
+        >>> patcher.patch(
+        ...     updates={"uid1": {"dataset": "imagenet", "epoch": 10}},
+        ...     parameters=session._get_view(),
+        ...     force=False,
+        ... )
+    """
+
+    def patch(
+        self,
+        *,
+        updates: dict[str, dict[str, Any]],
+        parameters: IParameterView,
+        force: bool = False,
+    ) -> None:
+        """Patch other_meta with conflict detection.
+
+        Args:
+            updates: Mapping of parameter uid -> {meta_key: value}.
+            parameters: Parameter view containing target parameters.
+            force: If True, overwrite existing metadata keys.
+
+        Raises:
+            MetadataPatchError: If unknown UIDs or conflicts are detected.
+        """
+        if not updates:
+            logger.debug("patch called with empty updates, skipping")
+            return
+
+        resolved, unknown_uids = self._resolve_parameters(updates, parameters)
+
+        if unknown_uids:
+            preview = ", ".join(sorted(unknown_uids)[:_ERROR_PREVIEW_LIMIT])
+            suffix = "..." if len(unknown_uids) > _ERROR_PREVIEW_LIMIT else ""
+            raise MetadataPatchError(
+                f"Unknown parameter UIDs (not found in session): {preview}{suffix}"
+            )
+
+        conflicts = self._validate_conflicts(updates, resolved, force)
+
+        if conflicts:
+            preview = ", ".join(
+                f"{uid}:{k}" for uid, k in conflicts[:_ERROR_PREVIEW_LIMIT]
+            )
+            suffix = "..." if len(conflicts) > _ERROR_PREVIEW_LIMIT else ""
+            raise MetadataPatchError(
+                f"Metadata conflicts (existing keys) for {len(conflicts)} entries: "
+                f"{preview}{suffix}"
+            )
+
+        self._apply_updates(updates, resolved, parameters)
+        logger.info("Patched metadata for %d parameters", len(updates))
+
+    def _resolve_parameters(
+        self,
+        updates: dict[str, dict[str, Any]],
+        parameters: IParameterView,
+    ) -> tuple[dict[str, ParameterDataProxy], list[str]]:
+        """Resolve parameter proxies from UIDs."""
+        resolved: dict[str, ParameterDataProxy] = {}
+        unknown_uids: list[str] = []
+
+        for uid in updates:
+            try:
+                resolved[uid] = parameters[uid]
+            except KeyError:
+                unknown_uids.append(uid)
+
+        return resolved, unknown_uids
+
+    def _validate_conflicts(
+        self,
+        updates: dict[str, dict[str, Any]],
+        resolved: dict[str, ParameterDataProxy],
+        force: bool,
+    ) -> list[tuple[str, str]]:
+        """Check for conflicting metadata keys."""
+        conflicts: list[tuple[str, str]] = []
+
+        for uid, meta_map in updates.items():
+            param = resolved[uid]
+            conflicts.extend(
+                (param.meta.uid, key)
+                for key in meta_map
+                if (not force) and (key in param.meta.other_meta)
+            )
+
+        return conflicts
+
+    def _apply_updates(
+        self,
+        updates: dict[str, dict[str, Any]],
+        resolved: dict[str, ParameterDataProxy],
+        parameters: IParameterView,
+    ) -> None:
+        """Apply metadata updates to parameters.
+
+        Updates other_meta in both the proxy and the metadata index.
+        """
+        total = len(updates)
+        repository = parameters._repository
+
+        with repository:
+            for uid, meta_map in tqdm(
+                updates.items(),
+                desc="Updating metadata...",
+                delay=PROGRESS_BAR_DELAY_SEC,
+                disable=total < PROGRESS_BAR_MIN_ITEMS,
+                total=total,
+            ):
+                param = resolved[uid]
+                old_meta = param.meta
+
+                # Build updated other_meta
+                new_other_meta = dict(old_meta.other_meta)
+                new_other_meta.update(meta_map)
+
+                # Create new metadata with updated other_meta
+                new_meta = replace(old_meta, other_meta=new_other_meta)
+
+                # Update the proxy's meta reference (frozen, so we use
+                # object.__setattr__)
+                object.__setattr__(param, "meta", new_meta)
+
+                new_meta_dict = new_meta.to_dict()
+                new_meta_dict.pop("uid")
+                repository.metadata_index.upsert(
+                    TABLE_PARAMETERS,
+                    uid=uid,
+                    **new_meta_dict,
+                )