diffract-core 0.2.1__py3-none-any.whl

diffract/session/namespaces/results/injester.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.aggregates.repository import AggregateRepository
+    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 FieldIngestionError(Exception):
+    """Error during field ingestion."""
+
+
+class FieldIngester:
+    """Utility for ingesting precomputed fields into parameters.
+
+    Validates and writes precomputed field values into parameters.
+
+    Example:
+        >>> ingester = FieldIngester()
+        >>> ingester.ingest(
+        ...     fields_by_uid={"param_uid": {"metric": 0.95, "loss": 0.05}},
+        ...     parameters=session._get_view(),
+        ...     force=False,
+        ... )
+    """
+
+    def ingest(
+        self,
+        *,
+        fields_by_uid: dict[str, dict[str, Any]],
+        parameters: IParameterView,
+        force: bool = False,
+    ) -> None:
+        """Ingest precomputed fields with validation.
+
+        Args:
+            fields_by_uid: Mapping of parameter uid -> {field_name: value}.
+            parameters: Parameter view containing target parameters.
+            force: If True, overwrite existing field values.
+
+        Raises:
+            FieldIngestionError: If unknown uids or conflicts are detected.
+        """
+        if not fields_by_uid:
+            logger.debug("ingest called with empty input, skipping")
+            return
+
+        resolved, unknown_uids = self._resolve_parameters(fields_by_uid, parameters)
+
+        if unknown_uids:
+            preview = ", ".join(sorted(unknown_uids)[:_ERROR_PREVIEW_LIMIT])
+            suffix = "..." if len(unknown_uids) > _ERROR_PREVIEW_LIMIT else ""
+            raise FieldIngestionError(
+                f"Unknown parameter UIDs (not found in session): {preview}{suffix}"
+            )
+
+        conflicts = self._validate_fields(fields_by_uid, resolved, force)
+
+        if conflicts:
+            preview = ", ".join(
+                f"{uid}:{f}" for uid, f in conflicts[:_ERROR_PREVIEW_LIMIT]
+            )
+            suffix = "..." if len(conflicts) > _ERROR_PREVIEW_LIMIT else ""
+            raise FieldIngestionError(
+                f"Field conflicts (already exist) for {len(conflicts)} entries: "
+                f"{preview}{suffix}"
+            )
+
+        self._write_fields(fields_by_uid, resolved, parameters)
+
+        total_fields = sum(len(fm) for fm in fields_by_uid.values())
+        logger.info(
+            "Ingested %d fields for %d parameters", total_fields, len(fields_by_uid)
+        )
+
+    def _resolve_parameters(
+        self,
+        fields_by_uid: 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 fields_by_uid:
+            try:
+                resolved[uid] = parameters[uid]
+            except KeyError:
+                unknown_uids.append(uid)
+
+        return resolved, unknown_uids
+
+    def _validate_fields(
+        self,
+        fields_by_uid: dict[str, dict[str, Any]],
+        resolved: dict[str, ParameterDataProxy],
+        force: bool,
+    ) -> list[tuple[str, str]]:
+        """Validate fields for conflict issues."""
+        conflicts: list[tuple[str, str]] = []
+
+        for uid, field_map in fields_by_uid.items():
+            param = resolved[uid]
+            conflicts.extend(
+                (uid, field_name)
+                for field_name in field_map
+                if (not force) and param.has_field(field_name)
+            )
+
+        return conflicts
+
+    def _write_fields(
+        self,
+        fields_by_uid: dict[str, dict[str, Any]],
+        resolved: dict[str, ParameterDataProxy],
+        parameters: IParameterView,
+    ) -> None:
+        """Write field values to parameters.
+
+        Note:
+            Uses batch write context from parameters view. If an error occurs
+            mid-write, partial writes may be committed depending on storage
+            backend. Callers should handle this as a potential partial failure.
+        """
+        with parameters:
+            for uid, field_map in fields_by_uid.items():
+                param = resolved[uid]
+                for field_name, value in field_map.items():
+                    param.set_field(field_name, value)
+
+
+class AggregateIngestionError(Exception):
+    """Error during aggregate ingestion."""
+
+
+class AggregateIngester:
+    """Utility for ingesting precomputed aggregate values.
+
+    Validates and writes precomputed aggregate values into the aggregate repository.
+
+    Example:
+        >>> ingester = AggregateIngester()
+        >>> ingester.ingest(
+        ...     aggregates=[
+        ...         {
+        ...             "field_name": "l_overlap",
+        ...             "context_models": ("m1", "m2"),
+        ...             "value": arr,
+        ...         }
+        ...     ],
+        ...     repository=aggregate_repository,
+        ...     force=False,
+        ... )
+    """
+
+    def ingest(
+        self,
+        *,
+        aggregates: list[dict[str, Any]],
+        repository: AggregateRepository,
+        force: bool = False,
+    ) -> None:
+        """Ingest precomputed aggregate values with validation.
+
+        Args:
+            aggregates: List of aggregate dictionaries, each containing:
+                - field_name: Base field name (e.g., "l_overlap").
+                - context_models: Tuple/list of model IDs.
+                - context_params: Tuple/list of parameter names (optional).
+                - value: The computed value to store.
+            repository: Target aggregate repository.
+            force: If True, overwrite existing aggregate values.
+
+        Raises:
+            AggregateIngestionError: If aggregate structure is invalid or
+                conflicts are detected (force=False).
+        """
+        if not aggregates:
+            logger.debug("ingest_aggregates called with empty input, skipping")
+            return
+
+        self._validate_structure(aggregates)
+
+        if not force:
+            conflicts = self._check_conflicts(aggregates, repository)
+            if conflicts:
+                preview = ", ".join(conflicts[:_ERROR_PREVIEW_LIMIT])
+                suffix = "..." if len(conflicts) > _ERROR_PREVIEW_LIMIT else ""
+                raise AggregateIngestionError(
+                    f"Aggregate conflicts detected: {preview}{suffix}"
+                )
+
+        self._write_aggregates(aggregates, repository)
+        logger.info("Ingested %d aggregates", len(aggregates))
+
+    def _validate_structure(self, aggregates: list[dict[str, Any]]) -> None:
+        """Validate aggregate structure."""
+        required_keys = ("field_name", "context_models", "value")
+
+        for i, agg in enumerate(aggregates):
+            for key in required_keys:
+                if key not in agg:
+                    raise AggregateIngestionError(
+                        f"Aggregate at index {i} missing '{key}'"
+                    )
+
+    def _check_conflicts(
+        self,
+        aggregates: list[dict[str, Any]],
+        repository: AggregateRepository,
+    ) -> list[str]:
+        """Check for existing aggregates that would conflict."""
+        from diffract.core.data.nn.aggregates.metadata import AggregateMetadata
+
+        existing_uids = set(repository.create_view().list_uids())
+        conflicts: list[str] = []
+
+        for agg in aggregates:
+            uid = AggregateMetadata.create_uid_from_context(
+                field_name=agg["field_name"],
+                context_models=tuple(agg["context_models"]),
+                context_params=tuple(agg.get("context_params", [])),
+            )
+            if uid in existing_uids:
+                conflicts.append(uid)
+
+        return conflicts
+
+    def _write_aggregates(
+        self,
+        aggregates: list[dict[str, Any]],
+        repository: AggregateRepository,
+    ) -> None:
+        """Write aggregate values to repository."""
+        with repository.storage_manager:
+            for agg in aggregates:
+                aggregate = repository.get_or_create(
+                    field_name=agg["field_name"],
+                    context_models=tuple(agg["context_models"]),
+                    context_params=tuple(agg.get("context_params", [])),
+                )
+                aggregate.set_field("value", agg["value"])

diffract/session/namespaces/utils/__init__.py

@@ -0,0 +1,141 @@
+"""Session utility namespace."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from dependency_injector.wiring import Provide, inject
+
+from diffract.session.namespaces.utils.merger import (
+    AggregateMerger,
+    MergeTargetState,
+    SessionMerger,
+)
+from diffract.session.session import Session, SessionContext
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.aggregates.repository import AggregateRepository
+    from diffract.core.data.nn.params.interface import IParameterRepository
+    from diffract.core.parallel import ParallelContext
+
+logger = logging.getLogger(__name__)
+
+
+class UtilsNamespace:
+    """Utility operations API for Session."""
+
+    @inject
+    def __init__(
+        self,
+        session_or_context: Session | SessionContext,
+        parameter_repository: IParameterRepository = Provide["nn.parameter_repository"],
+        aggregate_repository: AggregateRepository = Provide["nn.aggregate_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.__agg_repo = aggregate_repository
+        self.__parallel_context_factory = parallel_context_factory
+
+    def merge_other_session(
+        self,
+        other_session_or_context: Session | SessionContext,
+        fields: list[str] | None = None,
+        *,
+        verify: bool = True,
+        read_budget_bytes: int = 512 * 1024 * 1024,
+    ) -> None:
+        """Merge parameters from another session into this one.
+
+        Copies parameters and their computed fields from another session.
+        When verify is True, handles conflicts by skipping duplicate fields.
+
+        Args:
+            other_session_or_context: Source session to merge from.
+            fields: Specific fields to merge. If None, merges all fields.
+            verify: If True, check for conflicts and skip duplicates.
+            read_budget_bytes: Maximum bytes to read during merge operation.
+        """
+        with self.__session_or_context:
+            parallel = self.__parallel_context_factory()
+            target_state: MergeTargetState | None = None
+            if verify:
+                target_state = MergeTargetState()
+
+                models = self.__session_or_context.models
+                for parameter_info in models.parameters.list(verbose=True):
+                    model_id = parameter_info["model_id"]
+                    name = parameter_info["name"]
+                    uid = parameter_info["uid"]
+                    available_fields = parameter_info["available_fields"]
+
+                    if model_id not in target_state.models_mapping:
+                        target_state.models_mapping[model_id] = set()
+
+                    target_state.models_mapping[model_id].add(name)
+                    target_state.parameters_fields_mapping[(model_id, name)] = (
+                        available_fields
+                    )
+                    target_state.parameters_uids_mapping[(model_id, name)] = uid
+
+            if isinstance(other_session_or_context, Session):
+                source_context = other_session_or_context.filter()
+            else:
+                source_context = other_session_or_context
+
+            source_params = source_context._param_filter_context
+            source_aggs = source_context._agg_filter_context
+
+            # Merge parameters if any exist
+            if source_params:
+                logger.info(
+                    "Starting merge from session '%s' (%d parameters)",
+                    other_session_or_context,
+                    len(source_params),
+                )
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        "Merge source models: %s",
+                        other_session_or_context.models.list(),
+                    )
+
+                source_fields_by_uid = source_params.list_fields_by_uid(
+                    parallel=parallel
+                )
+
+                merger = SessionMerger(parallel=parallel)
+                merger.merge(
+                    source=source_params,
+                    source_fields_by_uid=source_fields_by_uid,
+                    target_repository=self.__param_repo,
+                    target_state=target_state,
+                    fields=fields,
+                    read_budget_bytes=read_budget_bytes,
+                )
+            else:
+                logger.debug("No parameters to merge from source session")
+
+            # Merge aggregates from source session (independent of parameters)
+            if source_aggs:
+                logger.info(
+                    "Starting merge from session '%s' (%d aggregates)",
+                    other_session_or_context,
+                    len(source_aggs),
+                )
+
+                aggregate_merger = AggregateMerger()
+                aggregate_merger.merge(
+                    source_aggregates=source_aggs,
+                    target_repository=self.__agg_repo,
+                    fields=fields,
+                    verify=verify,
+                )
+            else:
+                logger.debug("No aggregates to merge from source session")
+
+            # Invalidate field cache since parameters/fields may have been added
+            self.__session_or_context._field_cache.invalidate()

diffract/session/namespaces/utils/merger.py

@@ -0,0 +1,295 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from tqdm.auto import tqdm
+
+from diffract.core.constants import (
+    PROGRESS_BAR_DELAY_SEC,
+    PROGRESS_BAR_MIN_ITEMS,
+)
+from diffract.core.data.nn.params.proxy import ParameterDataProxy
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.aggregates.repository import AggregateRepository
+    from diffract.core.data.nn.aggregates.view import AggregateView
+    from diffract.core.data.nn.params.interface import (
+        IParameterProxy,
+        IParameterRepository,
+        IParameterView,
+    )
+    from diffract.core.parallel import ParallelContext
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MergeTargetState:
+    """State of target session for conflict detection during merge."""
+
+    models_mapping: dict[str, set[str]] = field(default_factory=dict)
+    parameters_uids_mapping: dict[tuple[str, str], str] = field(default_factory=dict)
+    parameters_fields_mapping: dict[tuple[str, str], list[str]] = field(
+        default_factory=dict
+    )
+
+
+class SessionMerger:
+    """Utility for merging parameters between sessions.
+
+    Handles chunked transfer of parameters with prefetching and conflict detection.
+
+    Example:
+        >>> merger = SessionMerger(parallel=parallel_context)
+        >>> merger.merge(
+        ...     source=source_session._get_view(),
+        ...     source_fields_by_uid=source.list_fields_by_uid(),
+        ...     target_repository=target_session._parameter_repository,
+        ...     target_state=target_session._build_merge_target_state(),
+        ...     fields=["metric", "loss"],
+        ... )
+    """
+
+    def __init__(self, parallel: ParallelContext) -> None:
+        """Initialize the merger.
+
+        Args:
+            parallel: Parallel context for concurrent operations.
+        """
+        self._parallel = parallel
+
+    def merge(
+        self,
+        *,
+        source: IParameterView,
+        source_fields_by_uid: dict[str, list[str]],
+        target_repository: IParameterRepository,
+        target_state: MergeTargetState | None = None,
+        fields: list[str] | None = None,
+        read_budget_bytes: int = 512 * 1024 * 1024,
+    ) -> None:
+        """Merge parameters from source view into target repository.
+
+        Args:
+            source: Source parameter view to merge from.
+            source_fields_by_uid: Precomputed mapping of uid -> available fields.
+            target_repository: Target repository to merge into.
+            target_state: Current state of target for conflict detection.
+                If None, no conflict checking is performed.
+            fields: Specific fields to merge. If None, merges all fields.
+            read_budget_bytes: Maximum bytes to read per chunk.
+        """
+        if not source:
+            return
+
+        total = len(source)
+        field_allowlist = set(fields) if fields else None
+
+        required_fields_by_uid = self._compute_required_fields(
+            source, source_fields_by_uid, field_allowlist
+        )
+
+        with tqdm(
+            total=total,
+            desc="Parameters migration...",
+            delay=PROGRESS_BAR_DELAY_SEC,
+            disable=total < PROGRESS_BAR_MIN_ITEMS,
+        ) as pbar:
+            for chunk in source.iter_chunks_by_read_budget(
+                budget_bytes=read_budget_bytes,
+                required_fields_by_uid=required_fields_by_uid,
+                parallel=self._parallel,
+            ):
+                self._process_chunk(
+                    chunk=chunk,
+                    source_fields_by_uid=source_fields_by_uid,
+                    target_repository=target_repository,
+                    target_state=target_state,
+                    field_allowlist=field_allowlist,
+                    required_fields_by_uid=required_fields_by_uid,
+                    pbar=pbar,
+                )
+
+    def _compute_required_fields(
+        self,
+        source: IParameterView,
+        source_fields_by_uid: dict[str, list[str]],
+        field_allowlist: set[str] | None,
+    ) -> dict[str, list[str]]:
+        """Compute required fields per parameter for prefetching."""
+        required_fields_by_uid: dict[str, list[str]] = {}
+
+        for param in source:
+            uid = param.meta.uid
+            available = source_fields_by_uid.get(uid, [])
+
+            required: set[str] = set()
+            for f in available:
+                if field_allowlist is not None and f not in field_allowlist:
+                    continue
+                required.add(f)
+
+            required_fields_by_uid[uid] = sorted(required)
+
+        return required_fields_by_uid
+
+    def _process_chunk(
+        self,
+        *,
+        chunk: IParameterView,
+        source_fields_by_uid: dict[str, list[str]],
+        target_repository: IParameterRepository,
+        target_state: MergeTargetState | None,
+        field_allowlist: set[str] | None,
+        required_fields_by_uid: dict[str, list[str]],
+        pbar: Any,
+    ) -> None:
+        """Process a single chunk of parameters during merge."""
+        chunk_required_fields_by_uid = {
+            p.meta.uid: required_fields_by_uid.get(p.meta.uid, []) for p in chunk
+        }
+        chunk.prefetch_fields(
+            fields_by_uid=chunk_required_fields_by_uid,
+            parallel=self._parallel,
+        )
+
+        with target_repository:
+            for pending_parameter in chunk:
+                self._transfer_parameter(
+                    pending_parameter=pending_parameter,
+                    source_fields_by_uid=source_fields_by_uid,
+                    target_repository=target_repository,
+                    target_state=target_state,
+                    field_allowlist=field_allowlist,
+                )
+                pbar.update(1)
+
+    def _transfer_parameter(
+        self,
+        *,
+        pending_parameter: IParameterProxy,
+        source_fields_by_uid: dict[str, list[str]],
+        target_repository: IParameterRepository,
+        target_state: MergeTargetState | None,
+        field_allowlist: set[str] | None,
+    ) -> None:
+        """Transfer a single parameter to the target repository."""
+        meta = pending_parameter.meta
+        create_new = True
+
+        if (
+            target_state is not None
+            and meta.model_id in target_state.models_mapping
+            and meta.name in target_state.models_mapping[meta.model_id]
+        ):
+            uid = target_state.parameters_uids_mapping[(meta.model_id, meta.name)]
+            new_parameter = target_repository.get_proxy(uid)
+            create_new = False
+
+        if create_new:
+            new_parameter = ParameterDataProxy.create_and_store(
+                meta,
+                repository=target_repository,
+            )
+
+        available_fields = source_fields_by_uid.get(meta.uid, [])
+        for field_name in available_fields:
+            if field_allowlist is not None and field_name not in field_allowlist:
+                continue
+            if (not create_new) and target_state is not None:
+                existing_fields = target_state.parameters_fields_mapping.get(
+                    (meta.model_id, meta.name), []
+                )
+                if field_name in existing_fields:
+                    continue
+
+            new_parameter.set_field(
+                field_name,
+                pending_parameter.get_field(field_name, auto_prefetch=True),
+            )
+
+
+class AggregateMerger:
+    """Utility for merging aggregates between sessions.
+
+    Handles transfer of aggregate entries with conflict detection.
+
+    Example:
+        >>> merger = AggregateMerger()
+        >>> merger.merge(
+        ...     source_aggregates=source_session._aggregate_repository.create_view(),
+        ...     target_repository=target_session._aggregate_repository,
+        ...     model_ids=["model_a"],
+        ...     fields=["l_overlap"],
+        ...     verify=True,
+        ... )
+    """
+
+    def merge(
+        self,
+        *,
+        source_aggregates: AggregateView,
+        target_repository: AggregateRepository,
+        model_ids: list[str] | None = None,
+        fields: list[str] | None = None,
+        verify: bool = True,
+    ) -> None:
+        """Merge aggregates from source view into target repository.
+
+        Args:
+            source_aggregates: Source aggregate view to merge from.
+            target_repository: Target repository to merge into.
+            model_ids: Filter source aggregates by context_models.
+            fields: Filter source aggregates by field_name.
+            verify: If True, skip existing aggregates.
+        """
+        filtered = source_aggregates
+        logger.debug("Merging aggregates: source has %d aggregates", len(filtered))
+
+        # Filter by model_ids if specified
+        if model_ids:
+            filtered = filtered.filter_by_context_models(*model_ids)
+
+        # Filter by fields if specified
+        if fields:
+            filtered = filtered.filter_by_field_name(*fields)
+
+        if not filtered:
+            logger.debug("No source aggregates to merge after filtering")
+            return
+
+        # Prefetch values from source
+        filtered.prefetch_fields(fields=["value"], parallel=None)
+
+        # Get existing aggregate UIDs in target
+        target_uids = set(target_repository.create_view().list_uids())
+
+        merged_count = 0
+        skipped_count = 0
+
+        with target_repository.storage_manager:
+            for source_agg in filtered:
+                if source_agg.meta.uid in target_uids and verify:
+                    skipped_count += 1
+                    continue
+
+                # Create or get aggregate in target
+                target_agg = target_repository.get_or_create(
+                    field_name=source_agg.meta.field_name,
+                    context_models=source_agg.meta.context_models,
+                    context_params=source_agg.meta.context_params,
+                )
+
+                # Copy value if exists
+                if source_agg.has_field("value"):
+                    target_agg.set_field("value", source_agg.get_field("value"))
+                    merged_count += 1
+
+        if merged_count > 0 or skipped_count > 0:
+            logger.debug(
+                "Aggregates merge: %d copied, %d skipped (existing)",
+                merged_count,
+                skipped_count,
+            )