arraybridge 0.2.0__py3-none-any.whl

arraybridge/gpu_cleanup.py

@@ -0,0 +1,149 @@
+"""
+GPU memory cleanup utilities for different frameworks.
+
+This module provides unified GPU memory cleanup functions for PyTorch, CuPy,
+TensorFlow, JAX, and pyclesperanto. The cleanup functions are designed to be called
+after processing steps to free up GPU memory that's no longer needed.
+
+REFACTORED: Uses enum-driven metaprogramming to eliminate 67% of code duplication.
+"""
+
+import gc
+import logging
+from typing import Optional
+
+from arraybridge.framework_config import _FRAMEWORK_CONFIG
+from arraybridge.types import MemoryType
+
+logger = logging.getLogger(__name__)
+
+
+
+
+
+
+
+def _create_cleanup_function(mem_type: MemoryType):
+    """
+    Factory function that creates a cleanup function for a specific memory type.
+
+    This single factory replaces 6 nearly-identical cleanup functions.
+    """
+    config = _FRAMEWORK_CONFIG[mem_type]
+    framework_name = config['import_name']
+    display_name = config['display_name']
+
+    # CPU memory type - no cleanup needed
+    if config['cleanup_ops'] is None:
+        def cleanup(device_id: Optional[int] = None) -> None:
+            """No-op cleanup for CPU memory type."""
+            logger.debug(f"🔥 GPU CLEANUP: No-op for {display_name} (CPU memory type)")
+
+        cleanup.__name__ = f"cleanup_{framework_name}_gpu"
+        cleanup.__doc__ = f"No-op cleanup for {display_name} (CPU memory type)."
+        return cleanup
+
+    # GPU memory type - generate cleanup function
+    def cleanup(device_id: Optional[int] = None) -> None:
+        """
+        Clean up {display_name} GPU memory.
+
+        Args:
+            device_id: Optional GPU device ID. If None, cleans all devices.
+        """
+        framework = globals().get(framework_name)
+
+        if framework is None:
+            logger.debug(f"{display_name} not available, skipping cleanup")
+            return
+
+        try:
+            # Check GPU availability
+            gpu_check_expr = config['gpu_check'].format(mod=framework_name)
+            try:
+                gpu_available = eval(gpu_check_expr, {framework_name: framework})
+            except Exception:
+                gpu_available = False
+
+            if not gpu_available:
+                return
+
+            # Execute cleanup operations
+            if device_id is not None and config['device_context'] is not None:
+                # Clean specific device with context
+                device_ctx_expr = config['device_context'].format(
+                    device_id=device_id, mod=framework_name
+                )
+                device_ctx = eval(device_ctx_expr, {framework_name: framework})
+
+                with device_ctx:
+                    # Execute cleanup operations
+                    cleanup_expr = config['cleanup_ops'].format(mod=framework_name)
+                    exec(cleanup_expr, {framework_name: framework, 'gc': gc})
+
+                logger.debug(f"🔥 GPU CLEANUP: Cleared {display_name} for device {device_id}")
+            else:
+                # Clean all devices (no device context)
+                cleanup_expr = config['cleanup_ops'].format(mod=framework_name)
+                exec(cleanup_expr, {framework_name: framework, 'gc': gc})
+                logger.debug(f"🔥 GPU CLEANUP: Cleared {display_name} for all devices")
+
+        except Exception as e:
+            logger.warning(f"Failed to cleanup {display_name} GPU memory: {e}")
+
+    # Set proper function name and docstring
+    cleanup.__name__ = f"cleanup_{framework_name}_gpu"
+    cleanup.__doc__ = cleanup.__doc__.format(display_name=display_name)
+
+    return cleanup
+
+
+# Auto-generate all cleanup functions
+for mem_type in MemoryType:
+    cleanup_func = _create_cleanup_function(mem_type)
+    globals()[cleanup_func.__name__] = cleanup_func
+
+
+# Auto-generate cleanup registry
+MEMORY_TYPE_CLEANUP_REGISTRY = {
+    mem_type.value: globals()[f"cleanup_{_FRAMEWORK_CONFIG[mem_type]['import_name']}_gpu"]
+    for mem_type in MemoryType
+}
+
+
+def cleanup_all_gpu_frameworks(device_id: Optional[int] = None) -> None:
+    """
+    Clean up GPU memory for all available frameworks.
+
+    This function calls cleanup for all GPU frameworks that are currently loaded.
+    It's safe to call even if some frameworks aren't available.
+
+    Args:
+        device_id: Optional GPU device ID. If None, cleans all devices.
+    """
+    logger.debug(f"🔥 GPU CLEANUP: Starting cleanup for all GPU frameworks (device_id={device_id})")
+
+    # Only cleanup GPU memory types (those with cleanup operations)
+    for mem_type, config in _FRAMEWORK_CONFIG.items():
+        if config['cleanup_ops'] is not None:
+            cleanup_func = MEMORY_TYPE_CLEANUP_REGISTRY[mem_type.value]
+            cleanup_func(device_id)
+
+    logger.debug("🔥 GPU CLEANUP: Completed cleanup for all GPU frameworks")
+
+
+
+
+
+# Export all cleanup functions and utilities
+__all__ = [
+    'cleanup_all_gpu_frameworks',
+    'MEMORY_TYPE_CLEANUP_REGISTRY',
+    'cleanup_numpy_gpu',  # noqa: F822
+    'cleanup_cupy_gpu',  # noqa: F822
+    'cleanup_torch_gpu',  # noqa: F822
+    'cleanup_tensorflow_gpu',  # noqa: F822
+    'cleanup_jax_gpu',  # noqa: F822
+    'cleanup_pyclesperanto_gpu',  # noqa: F822
+]
+

arraybridge/oom_recovery.py

@@ -0,0 +1,148 @@
+"""
+GPU Out of Memory (OOM) recovery utilities.
+
+Provides comprehensive OOM detection and cache clearing for all supported
+GPU frameworks in OpenHCS.
+
+REFACTORED: Uses enum-driven metaprogramming to eliminate 71% of code duplication.
+All OOM patterns and cache clearing operations are defined in framework_ops.py.
+"""
+
+import gc
+import logging
+from typing import Optional
+
+from arraybridge.framework_ops import _FRAMEWORK_OPS
+from arraybridge.types import MemoryType
+from arraybridge.utils import optional_import
+
+logger = logging.getLogger(__name__)
+
+
+def _is_oom_error(e: Exception, memory_type: str) -> bool:
+    """
+    Detect Out of Memory errors for all GPU frameworks.
+
+    Auto-generated from framework_ops.py OOM patterns.
+
+    Args:
+        e: Exception to check
+        memory_type: Memory type string (e.g., 'torch', 'cupy')
+
+    Returns:
+        True if exception is an OOM error for the given framework
+    """
+    # Find the MemoryType enum for this memory_type string
+    mem_type_enum = None
+    for mt in MemoryType:
+        if mt.value == memory_type:
+            mem_type_enum = mt
+            break
+
+    if mem_type_enum is None:
+        return False
+
+    ops = _FRAMEWORK_OPS[mem_type_enum]
+    error_str = str(e).lower()
+
+    # Check framework-specific exception types
+    for exc_type_expr in ops['oom_exception_types']:
+        try:
+            # Import the module and get the exception type
+            mod_name = ops['import_name']
+            mod = optional_import(mod_name)
+            if mod is None:
+                continue
+
+            # Evaluate the exception type expression
+            exc_type_str = exc_type_expr.format(mod='mod')
+            # Extract the attribute path
+            # (e.g., 'mod.cuda.OutOfMemoryError' -> ['cuda', 'OutOfMemoryError'])
+            parts = exc_type_str.split('.')[1:]  # Skip 'mod'
+            exc_type = mod
+            for part in parts:
+                if hasattr(exc_type, part):
+                    exc_type = getattr(exc_type, part)
+                else:
+                    exc_type = None
+                    break
+
+            if exc_type is not None and isinstance(e, exc_type):
+                return True
+        except Exception:
+            continue
+
+    # String-based detection using framework-specific patterns
+    return any(pattern in error_str for pattern in ops['oom_string_patterns'])
+
+
+def _clear_cache_for_memory_type(memory_type: str, device_id: Optional[int] = None):
+    """
+    Clear GPU cache for specific memory type.
+
+    Auto-generated from framework_ops.py cache clearing operations.
+
+    Args:
+        memory_type: Memory type string (e.g., 'torch', 'cupy')
+        device_id: GPU device ID (optional, currently unused but kept for API compatibility)
+    """
+    # Find the MemoryType enum for this memory_type string
+    mem_type_enum = None
+    for mt in MemoryType:
+        if mt.value == memory_type:
+            mem_type_enum = mt
+            break
+
+    if mem_type_enum is None:
+        logger.warning(f"Unknown memory type for cache clearing: {memory_type}")
+        gc.collect()
+        return
+
+    ops = _FRAMEWORK_OPS[mem_type_enum]
+
+    # Get the module
+    mod_name = ops['import_name']
+    mod = optional_import(mod_name)
+
+    if mod is None:
+        logger.warning(f"Module {mod_name} not available for cache clearing")
+        gc.collect()
+        return
+
+    # Execute cache clearing operations
+    cache_clear_expr = ops['oom_clear_cache']
+    if cache_clear_expr:
+        try:
+            # Execute cache clear directly (device context handled by the operations themselves)
+            exec(cache_clear_expr.format(mod=mod_name), {mod_name: mod, 'gc': gc})
+        except Exception as e:
+            logger.warning(f"Failed to clear cache for {memory_type}: {e}")
+
+    # Always trigger Python garbage collection
+    gc.collect()
+
+
+def _execute_with_oom_recovery(func_callable, memory_type: str, max_retries: int = 2):
+    """
+    Execute function with automatic OOM recovery.
+
+    Args:
+        func_callable: Function to execute
+        memory_type: Memory type from MemoryType enum
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        Function result
+
+    Raises:
+        Original exception if not OOM or retries exhausted
+    """
+    for attempt in range(max_retries + 1):
+        try:
+            return func_callable()
+        except Exception as e:
+            if not _is_oom_error(e, memory_type) or attempt == max_retries:
+                raise
+
+            # Clear cache and retry
+            _clear_cache_for_memory_type(memory_type)

arraybridge/slice_processing.py

@@ -0,0 +1,73 @@
+"""
+Shared slice-by-slice processing logic for all memory types.
+
+This module provides a single implementation of slice-by-slice processing
+that works for all memory types, eliminating duplication across dtype wrappers.
+"""
+
+from arraybridge.converters import detect_memory_type
+from arraybridge.stack_utils import stack_slices, unstack_slices
+
+
+def process_slices(image, func, args, kwargs):
+    """
+    Process a 3D array slice-by-slice using the provided function.
+
+    This function handles:
+    - Unstacking 3D arrays into 2D slices
+    - Processing each slice independently
+    - Handling functions that return tuples (main output + special outputs)
+    - Stacking results back into 3D arrays
+    - Combining special outputs from all slices
+
+    Args:
+        image: 3D array to process
+        func: Function to apply to each slice
+        args: Positional arguments to pass to func
+        kwargs: Keyword arguments to pass to func
+
+    Returns:
+        Processed 3D array, or tuple of (processed_3d_array, special_outputs...)
+        if func returns tuples
+    """
+    # Detect memory type and use proper OpenHCS utilities
+    memory_type = detect_memory_type(image)
+    gpu_id = 0  # Default GPU ID for slice processing
+
+    # Unstack 3D array into 2D slices
+    slices_2d = unstack_slices(image, memory_type, gpu_id)
+
+    # Process each slice and handle special outputs
+    main_outputs = []
+    special_outputs_list = []
+
+    for slice_2d in slices_2d:
+        slice_result = func(slice_2d, *args, **kwargs)
+
+        # Check if result is a tuple (indicating special outputs)
+        if isinstance(slice_result, tuple):
+            main_outputs.append(slice_result[0])  # First element is main output
+            special_outputs_list.append(slice_result[1:])  # Rest are special outputs
+        else:
+            main_outputs.append(slice_result)  # Single output
+
+    # Stack main outputs back into 3D array
+    result = stack_slices(main_outputs, memory_type, gpu_id)
+
+    # If we have special outputs, combine them and return tuple
+    if special_outputs_list:
+        # Combine special outputs from all slices
+        combined_special_outputs = []
+        num_special_outputs = len(special_outputs_list[0])
+
+        for i in range(num_special_outputs):
+            # Collect the i-th special output from all slices
+            special_output_values = [slice_outputs[i] for slice_outputs in special_outputs_list]
+            combined_special_outputs.append(special_output_values)
+
+        # Return tuple: (stacked_main_output, combined_special_output1,  # noqa: E501
+        # combined_special_output2, ...)
+        return (result, *combined_special_outputs)
+
+    return result
+

arraybridge/stack_utils.py

@@ -0,0 +1,317 @@
+"""
+Stack utilities module for OpenHCS.
+
+This module provides functions for stacking 2D slices into a 3D array
+and unstacking a 3D array into 2D slices, with explicit memory type handling.
+
+This module enforces Clause 278 — Mandatory 3D Output Enforcement:
+All functions must return a 3D array of shape [Z, Y, X], even when operating
+on a single 2D slice. No logic may check, coerce, or infer rank at unstack time.
+"""
+
+import logging
+from typing import Any
+
+from arraybridge.converters import detect_memory_type
+from arraybridge.framework_config import _FRAMEWORK_CONFIG
+from arraybridge.types import GPU_MEMORY_TYPES, MemoryType
+from arraybridge.utils import optional_import
+
+logger = logging.getLogger(__name__)
+
+# 🔍 MEMORY CONVERSION LOGGING: Test log to verify logger is working
+logger.debug("🔄 STACK_UTILS: Module loaded - memory conversion logging enabled")
+
+
+def _is_2d(data: Any) -> bool:
+    """
+    Check if data is a 2D array.
+
+    Args:
+        data: Data to check
+
+    Returns:
+        True if data is 2D, False otherwise
+    """
+    # Check if data has a shape attribute
+    if not hasattr(data, 'shape'):
+        return False
+
+    # Check if shape has length 2
+    return len(data.shape) == 2
+
+
+def _is_3d(data: Any) -> bool:
+    """
+    Check if data is a 3D array.
+
+    Args:
+        data: Data to check
+
+    Returns:
+        True if data is 3D, False otherwise
+    """
+    # Check if data has a shape attribute
+    if not hasattr(data, 'shape'):
+        return False
+
+    # Check if shape has length 3
+    return len(data.shape) == 3
+
+
+def _enforce_gpu_device_requirements(memory_type: str, gpu_id: int) -> None:
+    """
+    Enforce GPU device requirements.
+
+    Args:
+        memory_type: The memory type
+        gpu_id: The GPU device ID
+
+    Raises:
+        ValueError: If gpu_id is negative
+    """
+    # For GPU memory types, validate gpu_id
+    if memory_type in {mem_type.value for mem_type in GPU_MEMORY_TYPES}:
+        if gpu_id < 0:
+            raise ValueError(f"Invalid GPU device ID: {gpu_id}. Must be a non-negative integer.")
+
+
+# NOTE: Allocation operations now defined in framework_config.py
+# This eliminates the scattered _ALLOCATION_OPS dict
+
+
+def _allocate_stack_array(
+    memory_type: str, stack_shape: tuple, first_slice: Any, gpu_id: int
+) -> Any:
+    """
+    Allocate a 3D array for stacking slices using framework config.
+
+    Args:
+        memory_type: The target memory type
+        stack_shape: The shape of the stack (Z, Y, X)
+        first_slice: The first slice (used for dtype inference)
+        gpu_id: The GPU device ID
+
+    Returns:
+        Pre-allocated array or None for pyclesperanto
+    """
+    # Convert string to enum
+    mem_type = MemoryType(memory_type)
+    config = _FRAMEWORK_CONFIG[mem_type]
+    allocate_expr = config['allocate_stack']
+
+    # Check if allocation is None (pyclesperanto uses custom stacking)
+    if allocate_expr is None:
+        return None
+
+    # Import the module
+    mod = optional_import(mem_type.value)
+    if mod is None:
+        raise ValueError(f"{mem_type.value} is required for memory type {memory_type}")
+
+    # Handle dtype conversion if needed
+    needs_conversion = config['needs_dtype_conversion']
+    if callable(needs_conversion):
+        # It's a callable that determines if conversion is needed
+        needs_conversion = needs_conversion(first_slice, detect_memory_type)
+
+    if needs_conversion:
+        from arraybridge.converters import convert_memory
+        first_slice_source_type = detect_memory_type(first_slice)
+        sample_converted = convert_memory(  # noqa: F841 (used in eval)
+            data=first_slice,
+            source_type=first_slice_source_type,
+            target_type=memory_type,
+            gpu_id=gpu_id
+        )
+        dtype = sample_converted.dtype  # noqa: F841 (used in eval)
+    else:
+        dtype = first_slice.dtype if hasattr(first_slice, 'dtype') else None  # noqa: F841 (used in eval)
+
+    # Set up local variables for eval
+    np = optional_import("numpy")  # noqa: F841 (used in eval)
+    cupy = mod if mem_type == MemoryType.CUPY else None  # noqa: F841 (used in eval)
+    torch = mod if mem_type == MemoryType.TORCH else None  # noqa: F841 (used in eval)
+    tf = mod if mem_type == MemoryType.TENSORFLOW else None  # noqa: F841 (used in eval)
+    jnp = optional_import("jax.numpy") if mem_type == MemoryType.JAX else None  # noqa: F841 (used in eval)
+
+    # Execute allocation with context if needed
+    allocate_context = config.get('allocate_context')
+    if allocate_context:
+        context = eval(allocate_context)
+        with context:
+            return eval(allocate_expr)
+    else:
+        return eval(allocate_expr)
+
+
+def stack_slices(slices: list[Any], memory_type: str, gpu_id: int) -> Any:
+    """
+    Stack 2D slices into a 3D array with the specified memory type.
+
+    STRICT VALIDATION: Assumes all slices are 2D arrays.
+    No automatic handling of improper inputs.
+
+    Args:
+        slices: List of 2D slices (numpy arrays, cupy arrays, torch tensors, etc.)
+        memory_type: The memory type to use for the stacked array (REQUIRED)
+        gpu_id: The target GPU device ID (REQUIRED)
+
+    Returns:
+        A 3D array with the specified memory type of shape [Z, Y, X]
+
+    Raises:
+        ValueError: If memory_type is not supported or slices is empty
+        ValueError: If gpu_id is negative for GPU memory types
+        ValueError: If slices are not 2D arrays
+        MemoryConversionError: If conversion fails
+    """
+    if not slices:
+        raise ValueError("Cannot stack empty list of slices")
+
+    # Verify all slices are 2D
+    for i, slice_data in enumerate(slices):
+        if not _is_2d(slice_data):
+            raise ValueError(f"Slice at index {i} is not a 2D array. All slices must be 2D.")
+
+    # Analyze input types for conversion planning (minimal logging)
+    input_types = [detect_memory_type(slice_data) for slice_data in slices]
+    unique_input_types = set(input_types)
+    memory_type not in unique_input_types or len(unique_input_types) > 1
+
+    # Check GPU requirements
+    _enforce_gpu_device_requirements(memory_type, gpu_id)
+
+    # Pre-allocate the final 3D array to avoid intermediate list and final stack operation
+    first_slice = slices[0]
+    stack_shape = (len(slices), first_slice.shape[0], first_slice.shape[1])
+
+    # Create pre-allocated result array in target memory type using enum dispatch
+    result = _allocate_stack_array(memory_type, stack_shape, first_slice, gpu_id)
+
+    # Convert each slice and assign to result array
+    conversion_count = 0
+
+    # Check for custom stack handler (pyclesperanto)
+    mem_type = MemoryType(memory_type)
+    config = _FRAMEWORK_CONFIG[mem_type]
+    stack_handler = config.get('stack_handler')
+
+    if stack_handler:
+        # Use custom stack handler
+        mod = optional_import(mem_type.value)
+        result = stack_handler(slices, memory_type, gpu_id, mod)
+    else:
+        # Standard stacking logic
+        for i, slice_data in enumerate(slices):
+            source_type = detect_memory_type(slice_data)
+
+            # Track conversions for batch logging
+            if source_type != memory_type:
+                conversion_count += 1
+
+            # Direct conversion
+            if source_type == memory_type:
+                converted_data = slice_data
+            else:
+                from arraybridge.converters import convert_memory
+                converted_data = convert_memory(
+                    data=slice_data,
+                    source_type=source_type,
+                    target_type=memory_type,
+                    gpu_id=gpu_id
+                )
+
+            # Assign converted slice using framework-specific handler if available
+            assign_handler = config.get('assign_slice')
+            if assign_handler:
+                # Custom assignment (JAX immutability)
+                result = assign_handler(result, i, converted_data)
+            else:
+                # Standard assignment
+                result[i] = converted_data
+
+    # 🔍 MEMORY CONVERSION LOGGING: Only log when conversions happen or issues occur
+    if conversion_count > 0:
+        logger.debug(
+            f"🔄 STACK_SLICES: Converted {conversion_count}/{len(slices)} "
+            f"slices to {memory_type}"
+        )
+    # Silent success for no-conversion cases to reduce log pollution
+
+    return result
+
+
+def unstack_slices(
+    array: Any, memory_type: str, gpu_id: int, validate_slices: bool = True
+) -> list[Any]:
+    """
+    Split a 3D array into 2D slices along axis 0 and convert to the specified memory type.
+
+    STRICT VALIDATION: Input must be a 3D array. No automatic handling of improper inputs.
+
+    Args:
+        array: 3D array to split - MUST BE 3D
+        memory_type: The memory type to use for the output slices (REQUIRED)
+        gpu_id: The target GPU device ID (REQUIRED)
+        validate_slices: If True, validates that each extracted slice is 2D
+
+    Returns:
+        List of 2D slices in the specified memory type
+
+    Raises:
+        ValueError: If array is not 3D
+        ValueError: If validate_slices is True and any extracted slice is not 2D
+        ValueError: If gpu_id is negative for GPU memory types
+        ValueError: If memory_type is not supported
+        MemoryConversionError: If conversion fails
+    """
+    # Detect input type and check if conversion is needed
+    input_type = detect_memory_type(array)
+    getattr(array, 'shape', 'unknown')
+
+    # Verify the array is 3D - fail loudly if not
+    if not _is_3d(array):
+        raise ValueError(f"Array must be 3D, got shape {getattr(array, 'shape', 'unknown')}")
+
+    # Check GPU requirements
+    _enforce_gpu_device_requirements(memory_type, gpu_id)
+
+    # Convert to target memory type
+    source_type = input_type  # Reuse already detected type
+
+    # Direct conversion
+    if source_type == memory_type:
+        # No conversion needed - silent success to reduce log pollution
+        pass
+    else:
+        # Convert and log the conversion
+        from arraybridge.converters import convert_memory
+        logger.debug(f"🔄 UNSTACK_SLICES: Converting array - {source_type} → {memory_type}")
+        array = convert_memory(
+            data=array,
+            source_type=source_type,
+            target_type=memory_type,
+            gpu_id=gpu_id
+        )
+
+    # Extract slices along axis 0 (already in the target memory type)
+    slices = [array[i] for i in range(array.shape[0])]
+
+    # Validate that all extracted slices are 2D if requested
+    if validate_slices:
+        for i, slice_data in enumerate(slices):
+            if not _is_2d(slice_data):
+                raise ValueError(
+                    f"Extracted slice at index {i} is not 2D. "
+                    f"This indicates a malformed 3D array."
+                )
+
+    # 🔍 MEMORY CONVERSION LOGGING: Only log conversions or issues
+    if source_type != memory_type:
+        logger.debug(f"🔄 UNSTACK_SLICES: Converted and extracted {len(slices)} slices")
+    elif len(slices) == 0:
+        logger.warning("🔄 UNSTACK_SLICES: No slices extracted (empty array)")
+    # Silent success for no-conversion cases to reduce log pollution
+
+    return slices