diffract-core 0.2.1__py3-none-any.whl

diffract/core/parallel/execution.py

@@ -0,0 +1,91 @@
+"""Parallel execution utilities."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import time
+from collections.abc import Callable, Iterator, Sequence
+from itertools import chain
+from typing import TypeVar
+
+from diffract.core.utils.math import mean
+
+from .runtime import ParallelContext, should_parallelize
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def map_maybe_parallel(
+    items: Sequence[T],
+    fn: Callable[[T], R],
+    *,
+    parallel: ParallelContext | None,
+) -> Iterator[R]:
+    """Map fn over items with automatic parallelization decision.
+
+    Uses a small parallel pilot batch to estimate average task cost and
+    decide whether remaining work should be parallelized.
+
+    Args:
+        items: Sequence of items to process.
+        fn: Function to apply to each item.
+        parallel: Optional parallel context for execution.
+
+    Returns:
+        List of results from applying fn to each item.
+    """
+    if not items:
+        return iter([])
+
+    if parallel is None:
+        return map(fn, items)
+
+    executor = parallel.executor
+    workers = parallel.workers
+    overhead = parallel.calibration.submit_overhead_per_task_s
+
+    n = len(items)
+    if n <= 1 or workers <= 1:
+        return map(fn, items)
+
+    # Fast path: for small batches, avoid pilot overhead and preserve order.
+    if n <= workers:
+        return executor.map(fn, items)
+
+    # If overhead is already comparable to expected work, skip parallelism.
+    # This is a cheap guard for very small/fast tasks.
+    overhead = max(0.0, overhead)
+    pilot_n = min(n, workers)
+
+    def _timed(idx: int, item: T) -> tuple[int, float, R]:
+        t0 = time.perf_counter()
+        value = fn(item)
+        dt = max(0.0, time.perf_counter() - t0)
+        return idx, dt, value
+
+    futures = [executor.submit(_timed, i, items[i]) for i in range(pilot_n)]
+
+    pilot_results: list[R | None] = [None] * pilot_n
+    dts: list[float] = []
+    for fut in concurrent.futures.as_completed(futures):
+        idx, dt, value = fut.result()
+        pilot_results[idx] = value
+        dts.append(dt)
+
+    rest = items[pilot_n:]
+    if not rest:
+        return iter(pilot_results)
+
+    avg = mean(dts)
+    if should_parallelize(
+        avg_task_s=avg,
+        n_tasks=len(rest),
+        workers=workers,
+        overhead_per_task_s=overhead,
+    ):
+        rest_results = executor.map(fn, rest)
+    else:
+        rest_results = map(fn, rest)
+
+    return chain(pilot_results, rest_results)

diffract/core/parallel/runtime.py

@@ -0,0 +1,91 @@
+"""Runtime helpers for parallel execution decisions.
+
+This module intentionally does not depend on dependency-injector. It provides
+small, testable building blocks that can be wired via DI at the application
+level.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ParallelCalibration:
+    """Thread-pool overhead calibration data."""
+
+    workers: int
+    submit_overhead_per_task_s: float
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ParallelContext:
+    """Per-method parallel execution context."""
+
+    executor: Any
+    calibration: ParallelCalibration
+    workers: int
+
+
+def calibrate_thread_pool_overhead(
+    *, submit: Callable[[Callable[[int], int], int], object], workers: int
+) -> ParallelCalibration:
+    """Measure approximate per-task overhead for a running thread pool.
+
+    Args:
+        submit: Callable with signature like executor.submit(fn, arg).
+        workers: Configured worker count for the pool.
+
+    Returns:
+        Calibration object. The overhead is measured by submitting a small batch
+        of no-op tasks and waiting for completion.
+    """
+
+    def _noop(_: int) -> int:
+        return 0
+
+    n = max(1, workers * (workers + 1))
+    t0 = time.perf_counter()
+    futures = [submit(_noop, i) for i in range(n)]
+    for fut in futures:
+        # Executor API: future-like object with result()
+        _ = fut.result()
+    dt = max(0.0, time.perf_counter() - t0)
+    overhead = dt / n if n > 0 else 0.0
+    return ParallelCalibration(workers=workers, submit_overhead_per_task_s=overhead)
+
+
+@functools.lru_cache(maxsize=32)
+def get_thread_pool_calibration(workers: int) -> ParallelCalibration:
+    """Get cached overhead calibration for a worker count."""
+    import concurrent.futures
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+        return calibrate_thread_pool_overhead(
+            submit=executor.submit,
+            workers=workers,
+        )
+
+
+def should_parallelize(
+    *,
+    avg_task_s: float,
+    n_tasks: int,
+    workers: int,
+    overhead_per_task_s: float,
+) -> bool:
+    """Return True if parallel execution is expected to be faster."""
+    if n_tasks <= 1 or workers <= 1:
+        return False
+    avg = max(0.0, avg_task_s)
+    overhead = max(0.0, overhead_per_task_s)
+    sequential = avg * n_tasks
+    parallel_compute = (avg * n_tasks) / workers
+    parallel_total = parallel_compute + overhead * n_tasks
+    return sequential > parallel_total

diffract/core/storage/__init__.py

@@ -0,0 +1,35 @@
+"""Storage management module providing persistent data storage backends.
+
+This module offers a unified interface for persistent storage with support for
+multiple backends including HDF5. It uses dependency injection for flexible
+configuration and backend selection, providing durability guarantees for
+cached computations.
+
+Features:
+    - HDF5-based persistent storage with compression
+    - Hierarchical data organization by objects and fields
+    - Support for NumPy arrays, JSON-serializable data, and arbitrary objects
+    - Atomic write operations and data integrity
+    - Configurable compression and chunking strategies
+
+Example:
+    >>> from diffract.core.storage import StorageContainer
+    >>> container = StorageContainer()
+    >>> storage_manager = container.storage_manager()
+    >>> storage_manager.set_field("obj123", "result", numpy_array)
+"""
+
+from .containers import StorageContainer
+from .hdf5_manager import HDF5StorageManager
+from .interface import DEFAULT_TABLE, UID, IStorageManager
+from .ram_manager import RAMStorageManager
+from .zarr_manager import ZarrStorageManager
+
+__all__ = [
+    "DEFAULT_TABLE",
+    "UID",
+    "HDF5StorageManager",
+    "IStorageManager",
+    "StorageContainer",
+    "ZarrStorageManager",
+]

diffract/core/storage/base_manager.py

@@ -0,0 +1,330 @@
+import json
+import pickle
+import threading
+from abc import abstractmethod
+from contextlib import suppress
+from types import TracebackType
+from typing import Any, Self
+
+import numpy as np
+
+from diffract.core.storage.interface import DEFAULT_TABLE, UID, IStorageManager
+
+
+class BaseStorageManager(IStorageManager):
+    """Base storage manager with batching support and table separation."""
+
+    def __init__(
+        self,
+        batch_size_limit_bytes: int = 50 * 1024 * 1024,
+        batch_soft_limit_ratio: float = 0.9,
+        use_json: bool = True,
+    ) -> None:
+        self._batch_size_limit_bytes = batch_size_limit_bytes
+        self._batch_soft_limit_bytes = int(
+            batch_size_limit_bytes * batch_soft_limit_ratio
+        )
+        self._use_json = use_json
+
+        self._pending_set_bytes: int = 0
+        # Keys are (table, obj_uid, field_name)
+        self._set_field_sizes: dict[tuple[str, UID, str], int] = {}
+        self._set_field_batch: dict[tuple[str, UID, str], Any] = {}
+        self._erase_field_batch: set[tuple[str, UID, str]] = set()
+        self._erase_obj_batch: set[tuple[str, UID]] = set()
+        self._erase_field_for_all_batch: set[tuple[str, str]] = set()
+
+        self._thread_local = threading.local()  # Thread-local context depth
+        self._performing_batch_operations: bool = False
+
+    def has_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> bool:
+        """Check if field exists.
+
+        Args:
+            obj_uid: Object unique identifier.
+            field_name: Field name.
+            table: Table name for logical data separation.
+
+        Returns:
+            True if field exists, False otherwise.
+        """
+        self._auto_flush_before_read()
+        return self._has_field(obj_uid, field_name, table=table)
+
+    def get_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> Any:
+        """Retrieve field value.
+
+        Args:
+            obj_uid: Object unique identifier.
+            field_name: Field name.
+            table: Table name for logical data separation.
+
+        Returns:
+            Field value.
+
+        Raises:
+            KeyError: If field not found.
+        """
+        self._auto_flush_before_read()
+        return self._get_field(obj_uid, field_name, table=table)
+
+    def get_field_metadata(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> dict[str, Any] | None:
+        """Return stored metadata for a field if present."""
+        self._auto_flush_before_read()
+        return self._get_field_metadata(obj_uid, field_name, table=table)
+
+    def list_fields(
+        self, obj_uid: UID = None, *, table: str = DEFAULT_TABLE
+    ) -> list[str]:
+        """List fields.
+
+        Args:
+            obj_uid: Optional object unique identifier. If None, lists all fields.
+            table: Table name for logical data separation.
+
+        Returns:
+            List of field names.
+        """
+        self._auto_flush_before_read()
+        return self._list_fields(obj_uid, table=table)
+
+    def list_objs(self, *, table: str = DEFAULT_TABLE) -> list[str]:
+        """List all objects.
+
+        Args:
+            table: Table name for logical data separation.
+
+        Returns:
+            List of object unique identifiers.
+        """
+        self._auto_flush_before_read()
+        return self._list_objs(table=table)
+
+    def list_objs_has_field(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> list[UID]:
+        """List objects with specific field.
+
+        Args:
+            field_name: Field name.
+            table: Table name for logical data separation.
+
+        Returns:
+            List of object unique identifiers.
+        """
+        self._auto_flush_before_read()
+        return self._list_objs_has_field(field_name, table=table)
+
+    def set_field(
+        self, obj_uid: UID, field_name: str, value: Any, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Store field value.
+
+        Args:
+            obj_uid: Object unique identifier.
+            field_name: Field name.
+            value: Value to store.
+            table: Table name for logical data separation.
+        """
+        key = (table, obj_uid, field_name)
+        prev_size = self._set_field_sizes.pop(key, 0)
+        if prev_size:
+            self._pending_set_bytes = max(0, self._pending_set_bytes - prev_size)
+
+        estimated = self._estimate_value_size(value)
+        self._set_field_sizes[key] = estimated
+        self._pending_set_bytes += estimated
+        self._set_field_batch[key] = value
+
+        if not self._in_batch_context:
+            self._flush_set_field_batch()
+            return
+
+        if (
+            self._batch_size_limit_bytes > 0
+            and self._pending_set_bytes >= self._batch_soft_limit_bytes
+        ):
+            self._perform_batch_operations()
+
+    def erase_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove specific field.
+
+        Args:
+            obj_uid: Object unique identifier.
+            field_name: Field name.
+            table: Table name for logical data separation.
+        """
+        self._erase_field_batch.add((table, obj_uid, field_name))
+        if not self._in_batch_context:
+            self._flush_erase_field_batch()
+
+    def erase_obj(self, obj_uid: UID, *, table: str = DEFAULT_TABLE) -> None:
+        """Remove entire object and all its fields.
+
+        Args:
+            obj_uid: Object unique identifier.
+            table: Table name for logical data separation.
+        """
+        self._erase_obj_batch.add((table, obj_uid))
+        if not self._in_batch_context:
+            self._flush_erase_obj_batch()
+
+    def erase_field_for_all(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove field from all objects.
+
+        Args:
+            field_name: Field name.
+            table: Table name for logical data separation.
+        """
+        self._erase_field_for_all_batch.add((table, field_name))
+        if not self._in_batch_context:
+            self._flush_erase_field_for_all_batch()
+
+    def clear(self, *, table: str | None = None) -> None:
+        """Clear data from storage.
+
+        Args:
+            table: If provided, clear only this table. If None, clear all data.
+        """
+        self._set_field_batch.clear()
+        self._erase_field_batch.clear()
+        self._erase_obj_batch.clear()
+        self._erase_field_for_all_batch.clear()
+
+        self._clear(table=table)
+
+    def _auto_flush_before_read(self) -> None:
+        if self._in_batch_context and self._has_pending_operations:
+            self._perform_batch_operations()
+
+    @property
+    def _context_depth(self) -> int:
+        """Thread-local context depth for nested batch contexts."""
+        return getattr(self._thread_local, "context_depth", 0)
+
+    @_context_depth.setter
+    def _context_depth(self, value: int) -> None:
+        """Set thread-local context depth."""
+        self._thread_local.context_depth = value
+
+    def _estimate_value_size(self, value: Any) -> int:
+        """Best-effort estimate of serialized size in bytes."""
+        if isinstance(value, np.ndarray):
+            return int(value.nbytes)
+        if self._use_json:
+            with suppress(TypeError, ValueError):
+                return len(json.dumps(value, ensure_ascii=False).encode("utf-8"))
+        return len(pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL))
+
+    @property
+    def _in_batch_context(self) -> bool:
+        """True if inside a batch context (with statement)."""
+        return self._context_depth > 0
+
+    @property
+    def _has_pending_operations(self) -> bool:
+        """True if there are pending buffered operations that require flushing."""
+        return bool(
+            self._set_field_batch
+            or self._erase_field_batch
+            or self._erase_obj_batch
+            or self._erase_field_for_all_batch
+        )
+
+    def _perform_batch_operations(self) -> None:
+        """Execute all pending batch operations."""
+        if not self._performing_batch_operations:
+            self._performing_batch_operations = True
+            try:
+                self._flush_set_field_batch()
+                self._flush_erase_field_batch()
+                self._flush_erase_obj_batch()
+                self._flush_erase_field_for_all_batch()
+            finally:
+                self._performing_batch_operations = False
+
+    @abstractmethod
+    def _has_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> bool: ...
+
+    @abstractmethod
+    def _get_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> Any: ...
+
+    @abstractmethod
+    def _get_field_metadata(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> dict[str, Any] | None: ...
+
+    @abstractmethod
+    def _list_fields(
+        self, obj_uid: UID = None, *, table: str = DEFAULT_TABLE
+    ) -> list[str]: ...
+
+    @abstractmethod
+    def _list_objs(self, *, table: str = DEFAULT_TABLE) -> list[str]: ...
+
+    @abstractmethod
+    def _list_objs_has_field(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> list[UID]: ...
+
+    @abstractmethod
+    def _flush_set_field_batch(self) -> None: ...
+
+    @abstractmethod
+    def _flush_erase_field_batch(self) -> None: ...
+
+    @abstractmethod
+    def _flush_erase_obj_batch(self) -> None: ...
+
+    @abstractmethod
+    def _flush_erase_field_for_all_batch(self) -> None: ...
+
+    @abstractmethod
+    def _clear(self, *, table: str | None = None) -> None: ...
+
+    def __enter__(self) -> Self:
+        """Enter batch mode: queue operations and flush them on exit.
+
+        Supports nested context managers via reference counting. Only the
+        outermost entry clears batches; only the outermost exit flushes.
+        """
+        if self._context_depth == 0:
+            self._set_field_batch.clear()
+            self._set_field_sizes.clear()
+            self._pending_set_bytes = 0
+            self._erase_field_batch.clear()
+            self._erase_obj_batch.clear()
+            self._erase_field_for_all_batch.clear()
+        self._context_depth += 1
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit batch write context manager.
+
+        Only the outermost exit flushes operations and commits/rollbacks.
+        """
+        try:
+            if self._context_depth == 1:
+                self._perform_batch_operations()
+        finally:
+            self._context_depth = max(0, self._context_depth - 1)

diffract/core/storage/containers.py

@@ -0,0 +1,122 @@
+"""Dependency injection containers for storage management.
+
+This module provides dependency injection configuration for storage managers
+using the dependency-injector library. It supports multiple storage backends
+with configurable selection for persistent data storage.
+
+Example:
+    >>> container = StorageContainer()
+    >>> container.config.backend.from_value("hdf5")
+    >>> container.config.hdf5.path.from_value("/path/to/storage.h5")
+    >>> storage_manager = container.storage_manager()
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Generator
+from contextlib import contextmanager
+
+from dependency_injector import containers, providers
+
+from diffract.core.utils.build import build_with_defaults
+
+from .hdf5_manager import HDF5StorageManager
+from .hybrid_manager import HybridStorageManager
+from .interface import IStorageManager
+from .ram_manager import RAMStorageManager
+from .sqlite_manager import SQLiteStorageManager
+from .zarr_manager import ZarrStorageManager
+
+logger = logging.getLogger(__name__)
+
+
+class StorageContainer(containers.DeclarativeContainer):
+    """Dependency injection container for storage-related components.
+
+    Provides configurable storage manager selection between different
+    persistent storage backends. Currently supports HDF5 for hierarchical
+    data storage with compression and atomic operations.
+
+    The container manages storage backend lifecycle and configuration,
+    ensuring proper initialization and resource management.
+
+    Attributes:
+        config: Configuration provider for storage settings.
+        hdf5_manager: Singleton HDF5 storage manager provider.
+        storage_manager: Selector for active storage backend.
+    """
+
+    config = providers.Configuration()
+
+    hdf5_manager = providers.Singleton(
+        build_with_defaults,
+        HDF5StorageManager,
+        config.hdf5.as_(dict),
+    )
+
+    sqlite_manager = providers.Singleton(
+        build_with_defaults,
+        SQLiteStorageManager,
+        config.sqlite.as_(dict),
+    )
+
+    zarr_manager = providers.Singleton(
+        build_with_defaults,
+        ZarrStorageManager,
+        config.zarr.as_(dict),
+    )
+
+    ram_manager = providers.Singleton(RAMStorageManager)
+
+    # Selectors for hybrid storage backends
+    _light_backends = providers.Selector(
+        config.hybrid.light,
+        sqlite=sqlite_manager,
+        ram=ram_manager,
+        hdf5=hdf5_manager,
+        zarr=zarr_manager,
+    )
+
+    _heavy_backends = providers.Selector(
+        config.hybrid.heavy,
+        hdf5=hdf5_manager,
+        zarr=zarr_manager,
+        sqlite=sqlite_manager,
+        ram=ram_manager,
+    )
+
+    hybrid_manager = providers.Singleton(
+        HybridStorageManager,
+        light_storage=_light_backends,
+        heavy_storage=_heavy_backends,
+        array_threshold=config.hybrid.array_threshold,
+    )
+
+    storage_manager = providers.Selector(
+        config.backend,
+        hdf5=hdf5_manager,
+        ram=ram_manager,
+        sqlite=sqlite_manager,
+        hybrid=hybrid_manager,
+        zarr=zarr_manager,
+    )
+
+    storage_manager_resource = providers.Resource(
+        lambda manager: _connect_wrapper(manager),
+        storage_manager,
+    )
+
+
+@contextmanager
+def _connect_wrapper(manager: IStorageManager) -> Generator[None, None, None]:
+    logger.debug("Init resource: %s", manager.__class__.__name__)
+    if hasattr(manager, "connect"):
+        logger.debug("Init connection for: %s", manager.__class__.__name__)
+        manager.connect()
+    try:
+        yield manager
+    finally:
+        if hasattr(manager, "close"):
+            logger.debug("Close connection for: %s", manager.__class__.__name__)
+            manager.close()