oscura 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

oscura/core/memory_guard.py

@@ -0,0 +1,290 @@
+"""Memory-safe guards for TraceKit analysis.
+
+This module provides memory guards and resource limiting utilities to prevent
+out-of-memory conditions during analysis operations.
+
+
+Example:
+    >>> from oscura.core.memory_guard import MemoryGuard, check_memory_available
+    >>> if check_memory_available(500):  # Need at least 500MB
+    ...     with MemoryGuard(max_mb=1000, name="fft") as guard:
+    ...         result = compute_fft(data)
+    ...         if not guard.check():
+    ...             raise MemoryError("Exceeded memory limit")
+
+References:
+    See oscura.core.memory_monitor for runtime monitoring.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def get_memory_usage_mb() -> float:
+    """Get current process memory usage in MB.
+
+    Returns:
+        Current resident set size (RSS) in megabytes.
+
+    Example:
+        >>> mem_mb = get_memory_usage_mb()
+        >>> print(f"Current process using {mem_mb:.1f} MB")
+    """
+    try:
+        import psutil
+
+        process = psutil.Process(os.getpid())
+        return float(process.memory_info().rss / (1024 * 1024))
+    except ImportError:
+        # Fallback for systems without psutil
+        logger.debug("psutil not available, memory usage tracking disabled")
+        return 0.0
+
+
+def check_memory_available(required_mb: float = 100) -> bool:
+    """Check if sufficient memory is available.
+
+    Args:
+        required_mb: Required available memory in megabytes.
+
+    Returns:
+        True if sufficient memory is available.
+
+    Example:
+        >>> if not check_memory_available(1000):
+        ...     print("Warning: Less than 1GB available")
+        ...     # Reduce batch size or chunk operations
+    """
+    try:
+        import psutil
+
+        available = float(psutil.virtual_memory().available / (1024 * 1024))
+        return bool(available > required_mb)
+    except ImportError:
+        # Assume OK if we can't check
+        return True
+
+
+class MemoryGuard:
+    """Context manager for memory-safe operations.
+
+
+    Monitors memory usage within a context and raises warnings/errors
+    if limits are exceeded.
+
+    Attributes:
+        max_mb: Maximum memory limit in megabytes.
+        name: Operation name for logging.
+        start_mem: Starting memory usage (MB).
+
+    Example:
+        >>> with MemoryGuard(max_mb=2000, name="spectrogram") as guard:
+        ...     # Perform memory-intensive operation
+        ...     for chunk in data_chunks:
+        ...         process_chunk(chunk)
+        ...         if not guard.check():
+        ...             break  # Stop before exceeding limit
+        >>> stats = guard.get_stats()
+        >>> print(f"Peak: {stats['peak_mb']:.1f} MB, Delta: {stats['delta_mb']:.1f} MB")
+    """
+
+    def __init__(self, max_mb: float = 1000, name: str = "operation"):
+        """Initialize memory guard.
+
+        Args:
+            max_mb: Maximum memory increase allowed in megabytes.
+            name: Operation name for logging and error messages.
+        """
+        self.max_mb = max_mb
+        self.name = name
+        self.start_mem = 0.0
+        self._peak_mem = 0.0
+
+    def __enter__(self) -> MemoryGuard:
+        """Enter context and record starting memory."""
+        self.start_mem = get_memory_usage_mb()
+        self._peak_mem = self.start_mem
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> None:
+        """Exit context and report memory usage."""
+        # Note: exc_val and exc_tb intentionally unused but required for Python 3.11+ compatibility
+        end_mem = get_memory_usage_mb()
+        delta = end_mem - self.start_mem
+
+        if delta > self.max_mb:
+            logger.warning(
+                f"{self.name} used {delta:.1f} MB (limit: {self.max_mb:.1f} MB). "
+                f"Consider reducing batch size or enabling chunked processing."
+            )
+
+        # Update peak
+        self._peak_mem = max(self._peak_mem, end_mem)
+
+    def check(self) -> bool:
+        """Check if within memory limit.
+
+        Returns:
+            True if within limit, False if limit exceeded.
+
+        Example:
+            >>> with MemoryGuard(max_mb=500) as guard:
+            ...     for i in range(1000):
+            ...         # Do work
+            ...         if i % 100 == 0 and not guard.check():
+            ...             raise MemoryError("Memory limit exceeded")
+        """
+        current = get_memory_usage_mb()
+        self._peak_mem = max(self._peak_mem, current)
+        delta = current - self.start_mem
+
+        if delta > self.max_mb:
+            logger.warning(
+                f"{self.name}: Memory usage {delta:.1f} MB exceeds limit {self.max_mb:.1f} MB"
+            )
+            return False
+
+        return True
+
+    def get_stats(self) -> dict[str, float]:
+        """Get memory statistics for this guard.
+
+        Returns:
+            Dictionary with keys:
+                - start_mb: Starting memory
+                - current_mb: Current memory
+                - peak_mb: Peak memory
+                - delta_mb: Memory increase since start
+                - limit_mb: Configured limit
+
+        Example:
+            >>> with MemoryGuard(max_mb=1000, name="test") as guard:
+            ...     # ... work ...
+            ...     pass
+            >>> stats = guard.get_stats()
+            >>> print(f"Used {stats['delta_mb']:.1f} / {stats['limit_mb']:.1f} MB")
+        """
+        current = get_memory_usage_mb()
+        return {
+            "start_mb": self.start_mem,
+            "current_mb": current,
+            "peak_mb": self._peak_mem,
+            "delta_mb": current - self.start_mem,
+            "limit_mb": self.max_mb,
+        }
+
+
+def safe_array_size(shape: tuple[int, ...], dtype_bytes: int = 8) -> int:
+    """Calculate array size in bytes, checking for overflow.
+
+
+    Args:
+        shape: Array shape tuple.
+        dtype_bytes: Bytes per element (default: 8 for float64).
+
+    Returns:
+        Total array size in bytes.
+
+    Raises:
+        OverflowError: If array size would overflow.
+
+    Example:
+        >>> size = safe_array_size((1000, 1000, 8), dtype_bytes=8)
+        >>> print(f"Array would use {size / 1e6:.1f} MB")
+        >>> # Check if safe to allocate
+        >>> if can_allocate(size):
+        ...     arr = np.zeros((1000, 1000, 8))
+    """
+    try:
+        import numpy as np
+
+        total_elements = np.prod(shape)
+
+        # Check for overflow in element count
+        if total_elements > sys.maxsize // dtype_bytes:
+            raise OverflowError(f"Array size too large: {shape}")
+
+        size = int(total_elements) * dtype_bytes
+        return size
+
+    except (OverflowError, ValueError) as e:
+        raise OverflowError(f"Array dimensions {shape} would cause overflow") from e
+
+
+def can_allocate(size_bytes: int) -> bool:
+    """Check if allocation is safe given available memory.
+
+
+    Args:
+        size_bytes: Requested allocation size in bytes.
+
+    Returns:
+        True if allocation is safe (with 2x safety margin).
+
+    Example:
+        >>> import numpy as np
+        >>> shape = (10000, 10000)
+        >>> size = safe_array_size(shape, dtype_bytes=8)
+        >>> if can_allocate(size):
+        ...     arr = np.zeros(shape)
+        ... else:
+        ...     print("Not enough memory, use chunked processing")
+    """
+    size_mb = size_bytes / (1024 * 1024)
+
+    # Check with 2x safety margin
+    return check_memory_available(size_mb * 2)
+
+
+def get_safe_chunk_size(
+    total_samples: int,
+    dtype_bytes: int = 8,
+    max_chunk_mb: float = 100,
+) -> int:
+    """Calculate safe chunk size for processing large datasets.
+
+
+    Args:
+        total_samples: Total number of samples to process.
+        dtype_bytes: Bytes per sample (default: 8 for float64).
+        max_chunk_mb: Maximum chunk size in megabytes.
+
+    Returns:
+        Chunk size in samples that fits within memory limit.
+
+    Example:
+        >>> total = 1_000_000_000  # 1 billion samples
+        >>> chunk_size = get_safe_chunk_size(total, max_chunk_mb=100)
+        >>> print(f"Process in chunks of {chunk_size:,} samples")
+        >>> for i in range(0, total, chunk_size):
+        ...     chunk = data[i:i+chunk_size]
+        ...     process(chunk)
+    """
+    max_bytes = max_chunk_mb * 1024 * 1024
+    max_samples = max_bytes // dtype_bytes
+
+    # Ensure at least 1000 samples per chunk, but not more than total
+    chunk_size = max(1000, min(max_samples, total_samples))
+
+    return int(chunk_size)
+
+
+__all__ = [
+    "MemoryGuard",
+    "can_allocate",
+    "check_memory_available",
+    "get_memory_usage_mb",
+    "get_safe_chunk_size",
+    "safe_array_size",
+]

oscura/core/memory_limits.py

@@ -0,0 +1,336 @@
+"""Per-operation memory limits for TraceKit.
+
+This module provides fine-grained memory control for individual operations
+with automatic parameter adjustment to fit memory constraints.
+
+
+Example:
+    >>> from oscura.core.memory_limits import apply_memory_limit
+    >>> params = apply_memory_limit('spectrogram', samples=1e9, max_memory='512MB')
+    >>> print(f"Adjusted nperseg: {params['nperseg']}")
+
+References:
+    See oscura.config.memory for global memory configuration.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+from oscura.config.memory import get_memory_config
+from oscura.utils.memory import estimate_memory
+
+
+def parse_memory_limit(limit: int | str | None) -> int | None:
+    """Parse memory limit from various formats.
+
+    Args:
+        limit: Memory limit as bytes (int), string ("4GB", "512MB"), or None.
+
+    Returns:
+        Memory limit in bytes, or None for no limit.
+
+    Raises:
+        ValueError: If format is invalid.
+
+    Example:
+        >>> parse_memory_limit("4GB")
+        4000000000
+        >>> parse_memory_limit(512 * 1024**2)
+        536870912
+        >>> parse_memory_limit(None) is None
+        True
+    """
+    if limit is None:
+        return None
+
+    if isinstance(limit, str):
+        limit_upper = limit.upper().strip()
+        try:
+            if limit_upper.endswith("GB"):
+                return int(float(limit_upper[:-2]) * 1e9)
+            elif limit_upper.endswith("MB"):
+                return int(float(limit_upper[:-2]) * 1e6)
+            elif limit_upper.endswith("KB"):
+                return int(float(limit_upper[:-2]) * 1e3)
+            elif limit_upper.endswith("GIB"):
+                return int(float(limit_upper[:-3]) * 1024**3)
+            elif limit_upper.endswith("MIB"):
+                return int(float(limit_upper[:-3]) * 1024**2)
+            elif limit_upper.endswith("KIB"):
+                return int(float(limit_upper[:-3]) * 1024)
+            else:
+                return int(float(limit_upper))
+        except ValueError as e:
+            raise ValueError(f"Invalid memory limit format: {limit}") from e
+
+    return int(limit)
+
+
+def apply_memory_limit(
+    operation: str,
+    samples: int | float,
+    *,
+    max_memory: int | str | None = None,
+    **params: Any,
+) -> dict[str, Any]:
+    """Apply memory limit and adjust parameters to fit.
+
+
+    Args:
+        operation: Operation name (fft, psd, spectrogram, etc.).
+        samples: Number of samples to process.
+        max_memory: Maximum memory limit (overrides global config if provided).
+        **params: Operation parameters to adjust.
+
+    Returns:
+        Adjusted parameters dictionary that fits within memory limit.
+
+    Example:
+        >>> params = apply_memory_limit('spectrogram', samples=1e9, max_memory='512MB', nperseg=8192)
+        >>> print(f"Adjusted to nperseg={params['nperseg']} to fit 512MB")
+
+    Note:
+        If parameters cannot be adjusted to fit memory, a warning is issued
+        and the original parameters are returned.
+    """
+    # Parse memory limit
+    limit_bytes = parse_memory_limit(max_memory)
+    if limit_bytes is None:
+        # Use global config
+        config = get_memory_config()
+        limit_bytes = config.max_memory
+        if limit_bytes is None:
+            # No limit, return params unchanged
+            return params
+
+    samples = int(samples)
+
+    # Estimate with current parameters
+    current_estimate = estimate_memory(operation, samples, **params)
+
+    if current_estimate.total <= limit_bytes:
+        # Already within limit
+        return params
+
+    # Need to adjust parameters
+    adjusted_params = params.copy()
+
+    if operation in ("fft", "psd"):
+        # Reduce nfft if specified
+        if "nfft" in adjusted_params:
+            # Try reducing nfft
+            original_nfft = adjusted_params["nfft"]
+            # Binary search for suitable nfft
+            nfft = _find_max_nfft(operation, samples, limit_bytes, **adjusted_params)
+            if nfft < original_nfft:
+                adjusted_params["nfft"] = nfft
+                warnings.warn(
+                    f"Reduced nfft from {original_nfft} to {nfft} to fit {limit_bytes / 1e6:.1f} MB limit",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+    elif operation == "spectrogram":
+        # Adjust nperseg and/or nfft
+        original_nperseg = adjusted_params.get("nperseg", 256)
+        original_nfft = adjusted_params.get("nfft", original_nperseg)
+
+        # Try reducing nperseg first
+        nperseg = _find_max_nperseg(samples, limit_bytes, noverlap=adjusted_params.get("noverlap"))
+        if nperseg < original_nperseg:
+            adjusted_params["nperseg"] = nperseg
+            # Adjust noverlap proportionally
+            if "noverlap" in adjusted_params:
+                overlap_ratio = adjusted_params["noverlap"] / original_nperseg
+                adjusted_params["noverlap"] = int(nperseg * overlap_ratio)
+            warnings.warn(
+                f"Reduced nperseg from {original_nperseg} to {nperseg} to fit {limit_bytes / 1e6:.1f} MB limit",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        # Also reduce nfft if needed
+        if "nfft" in adjusted_params and adjusted_params["nfft"] > nperseg:
+            adjusted_params["nfft"] = nperseg
+
+    elif operation == "eye_diagram":
+        # Reduce samples_per_ui or num_uis
+        if "num_uis" in adjusted_params:
+            original_num_uis = adjusted_params["num_uis"]
+            # Calculate max num_uis that fits
+            samples_per_ui = adjusted_params.get("samples_per_ui", 100)
+            max_num_uis = _find_max_num_uis(limit_bytes, samples_per_ui)
+            if max_num_uis < original_num_uis:
+                adjusted_params["num_uis"] = max_num_uis
+                warnings.warn(
+                    f"Reduced num_uis from {original_num_uis} to {max_num_uis} to fit {limit_bytes / 1e6:.1f} MB limit",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+    # Verify final estimate
+    final_estimate = estimate_memory(operation, samples, **adjusted_params)
+    if final_estimate.total > limit_bytes:
+        warnings.warn(
+            f"Could not adjust parameters to fit {limit_bytes / 1e6:.1f} MB limit. "
+            f"Operation requires {final_estimate.total / 1e6:.1f} MB. "
+            "Consider using chunked processing or increasing memory limit.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    return adjusted_params
+
+
+def _find_max_nfft(operation: str, samples: int, limit_bytes: int, **params: Any) -> int:
+    """Binary search for maximum nfft that fits memory limit.
+
+    Args:
+        operation: Operation name.
+        samples: Number of samples.
+        limit_bytes: Memory limit in bytes.
+        **params: Additional parameters.
+
+    Returns:
+        Maximum nfft that fits within limit.
+    """
+    min_nfft = 64
+    max_nfft = params.get("nfft", 8192)
+
+    # Binary search
+    while min_nfft < max_nfft:
+        mid_nfft = (min_nfft + max_nfft + 1) // 2
+        test_params = {**params, "nfft": mid_nfft}
+        estimate = estimate_memory(operation, samples, **test_params)
+
+        if estimate.total <= limit_bytes:
+            min_nfft = mid_nfft
+        else:
+            max_nfft = mid_nfft - 1
+
+    return min_nfft
+
+
+def _find_max_nperseg(samples: int, limit_bytes: int, noverlap: int | None = None) -> int:
+    """Binary search for maximum nperseg that fits memory limit.
+
+    Args:
+        samples: Number of samples.
+        limit_bytes: Memory limit in bytes.
+        noverlap: Overlap samples (if specified).
+
+    Returns:
+        Maximum nperseg that fits within limit.
+    """
+    min_nperseg = 64
+    max_nperseg = min(8192, samples // 4)
+
+    # Binary search
+    while min_nperseg < max_nperseg:
+        mid_nperseg = (min_nperseg + max_nperseg + 1) // 2
+
+        # Calculate memory for this nperseg
+        hop = mid_nperseg - (noverlap or mid_nperseg // 2)
+        num_segments = max(1, (samples - (noverlap or mid_nperseg // 2)) // hop)
+
+        # Estimate memory
+        bytes_per_sample = 8  # float64
+        data_mem = samples * bytes_per_sample
+        intermediate_mem = mid_nperseg * bytes_per_sample * 2 + mid_nperseg * bytes_per_sample * 2
+        output_mem = (mid_nperseg // 2 + 1) * num_segments * bytes_per_sample * 2
+
+        total_mem = data_mem + intermediate_mem + output_mem
+
+        if total_mem <= limit_bytes:
+            min_nperseg = mid_nperseg
+        else:
+            max_nperseg = mid_nperseg - 1
+
+    return min_nperseg
+
+
+def _find_max_num_uis(limit_bytes: int, samples_per_ui: int) -> int:
+    """Find maximum num_uis that fits memory limit for eye diagrams.
+
+    Args:
+        limit_bytes: Memory limit in bytes.
+        samples_per_ui: Samples per unit interval.
+
+    Returns:
+        Maximum num_uis that fits.
+    """
+    bytes_per_sample = 8  # float64
+    # Eye diagram memory: samples_per_ui * num_uis * bytes_per_sample
+    max_num_uis = limit_bytes // (samples_per_ui * bytes_per_sample * 2)
+    return max(100, int(max_num_uis))  # At least 100 UIs
+
+
+def get_operation_memory_limit(
+    operation: str,
+    max_memory: int | str | None = None,
+) -> int:
+    """Get effective memory limit for an operation.
+
+    Args:
+        operation: Operation name.
+        max_memory: Override limit (or None for global config).
+
+    Returns:
+        Memory limit in bytes.
+
+    Example:
+        >>> limit = get_operation_memory_limit('spectrogram', max_memory='512MB')
+        >>> print(f"Limit: {limit / 1e6:.1f} MB")
+    """
+    # Parse override
+    limit_bytes = parse_memory_limit(max_memory)
+    if limit_bytes is not None:
+        return limit_bytes
+
+    # Use global config
+    config = get_memory_config()
+    if config.max_memory is not None:
+        return config.max_memory
+
+    # Default: 80% of available
+    from oscura.utils.memory import get_available_memory
+
+    return int(get_available_memory() * 0.8)
+
+
+def check_operation_fits(
+    operation: str,
+    samples: int | float,
+    *,
+    max_memory: int | str | None = None,
+    **params: Any,
+) -> bool:
+    """Check if operation with given parameters fits within memory limit.
+
+    Args:
+        operation: Operation name.
+        samples: Number of samples.
+        max_memory: Memory limit (or None for global config).
+        **params: Operation parameters.
+
+    Returns:
+        True if operation fits within limit.
+
+    Example:
+        >>> fits = check_operation_fits('fft', samples=1e9, max_memory='4GB', nfft=8192)
+        >>> if not fits:
+        ...     print("FFT too large for 4GB limit")
+    """
+    limit_bytes = get_operation_memory_limit(operation, max_memory)
+    estimate = estimate_memory(operation, samples, **params)
+    return estimate.total <= limit_bytes
+
+
+__all__ = [
+    "apply_memory_limit",
+    "check_operation_fits",
+    "get_operation_memory_limit",
+    "parse_memory_limit",
+]