Sync OpenHCS io

Author: trissim

src/polystore/__init__.py

@@ -1,85 +1,140 @@
 """
-Polystore - Framework-agnostic multi-backend storage abstraction.
-
-Provides pluggable storage backends with multi-framework I/O support for
-NumPy, PyTorch, JAX, TensorFlow, CuPy, and Zarr.
+Polystore package exports.
 """
 
-__version__ = "0.1.2"
+import os
 
-# Core abstractions
+from .atomic import file_lock, atomic_write_json, atomic_update_json, FileLockError, FileLockTimeoutError
+from .backend_registry import (
+    get_backend_instance,
+    cleanup_backend_connections,
+    cleanup_all_backends,
+    register_cleanup_callback,
+    STORAGE_BACKENDS,
+)
 from .base import (
+    BackendBase,
     DataSink,
     DataSource,
-    StorageBackend,
-    VirtualBackend,
     ReadOnlyBackend,
+    StorageBackend,
+    storage_registry,
+    reset_memory_backend,
+    ensure_storage_registry,
+    get_backend,
 )
-
-# Concrete backends
-from .memory import MemoryBackend
-from .disk import DiskBackend
-
-# Optional backends
-# Zarr is a required backend for this project. Import it directly so
-# missing dependency errors surface loudly during test/setup.
-from .zarr import ZarrBackend
-
-# File manager
+from .constants import Backend, MemoryType, TransportMode
+from .disk import DiskStorageBackend
 from .filemanager import FileManager
-
-# Registry
-from .backend_registry import BackendRegistry, create_storage_registry
-
-# Atomic operations
-from .atomic import atomic_write, atomic_write_json
-
-# Exceptions
-from .exceptions import (
-    StorageError,
-    StorageResolutionError,
-    BackendNotFoundError,
-    UnsupportedFormatError,
+from .formats import FileFormat, DEFAULT_IMAGE_EXTENSIONS
+from .memory import MemoryStorageBackend
+from .metadata_writer import (
+    AtomicMetadataWriter,
+    MetadataWriteError,
+    METADATA_CONFIG,
+    get_metadata_path,
+    get_subdirectory_name,
+    resolve_subdirectory_path,
 )
-
-# Streaming (optional and lazy)
-# Don't import at module level - streaming is heavy and optional
-# Users can import manually if needed: from polystore.streaming import StreamingBackend
-try:
-    from .streaming import StreamingBackend
-except ImportError:
-    StreamingBackend = None
-# Streaming (optional and lazy)
-# Don't import at module level - streaming is heavy and optional
-# Users can import manually if needed: from polystore.streaming import StreamingBackend
-StreamingBackend = None
+from .metadata_migration import detect_legacy_format, migrate_legacy_metadata, migrate_plate_metadata
+from .pipeline_migration import detect_legacy_pipeline, migrate_pipeline_file, load_pipeline_with_migration
+from .roi import (
+    ROI,
+    PolygonShape,
+    PolylineShape,
+    MaskShape,
+    PointShape,
+    EllipseShape,
+    extract_rois_from_labeled_mask,
+    load_rois_from_json,
+    load_rois_from_zip,
+    materialize_rois,
+)
+from .streaming import StreamingBackend
+from .streaming_constants import StreamingDataType, NapariShapeType
 
 __all__ = [
-    # Version
-    "__version__",
-    # Core abstractions
+    "Backend",
+    "MemoryType",
+    "TransportMode",
+    "FileFormat",
+    "DEFAULT_IMAGE_EXTENSIONS",
+    "BackendBase",
     "DataSink",
     "DataSource",
-    "StorageBackend",
-    "VirtualBackend",
     "ReadOnlyBackend",
-    # Backends
-    "MemoryBackend",
-    "DiskBackend",
-    "ZarrBackend",
-    # File manager
+    "StorageBackend",
+    "StreamingBackend",
+    "storage_registry",
+    "reset_memory_backend",
+    "ensure_storage_registry",
+    "get_backend",
+    "get_backend_instance",
+    "cleanup_backend_connections",
+    "cleanup_all_backends",
+    "register_cleanup_callback",
+    "STORAGE_BACKENDS",
+    "DiskStorageBackend",
+    "MemoryStorageBackend",
     "FileManager",
-    # Registry
-    "BackendRegistry",
-    # Atomic operations
-    "atomic_write",
+    "file_lock",
     "atomic_write_json",
-    # Exceptions
-    "StorageError",
-    "StorageResolutionError",
-    "BackendNotFoundError",
-    "UnsupportedFormatError",
-    # Streaming (optional)
-    "StreamingBackend",
+    "atomic_update_json",
+    "FileLockError",
+    "FileLockTimeoutError",
+    "AtomicMetadataWriter",
+    "MetadataWriteError",
+    "METADATA_CONFIG",
+    "get_metadata_path",
+    "get_subdirectory_name",
+    "resolve_subdirectory_path",
+    "detect_legacy_format",
+    "migrate_legacy_metadata",
+    "migrate_plate_metadata",
+    "detect_legacy_pipeline",
+    "migrate_pipeline_file",
+    "load_pipeline_with_migration",
+    "ROI",
+    "PolygonShape",
+    "PolylineShape",
+    "MaskShape",
+    "PointShape",
+    "EllipseShape",
+    "extract_rois_from_labeled_mask",
+    "load_rois_from_json",
+    "load_rois_from_zip",
+    "materialize_rois",
+    "StreamingDataType",
+    "NapariShapeType",
+    "NapariStreamingBackend",
+    "FijiStreamingBackend",
+    "ZarrStorageBackend",
+    "OMEROLocalBackend",
+    "OMEROFileFormatRegistry",
 ]
 
+_LAZY_BACKEND_REGISTRY = {
+    "NapariStreamingBackend": ("polystore.napari_stream", "NapariStreamingBackend"),
+    "FijiStreamingBackend": ("polystore.fiji_stream", "FijiStreamingBackend"),
+    "ZarrStorageBackend": ("polystore.zarr", "ZarrStorageBackend"),
+    "OMEROLocalBackend": ("polystore.omero_local", "OMEROLocalBackend"),
+    "OMEROFileFormatRegistry": ("polystore.omero_local", "OMEROFileFormatRegistry"),
+}
+
+
+def __getattr__(name):
+    """Lazy import of optional/extra backend classes."""
+    if name not in _LAZY_BACKEND_REGISTRY:
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+    if os.getenv("POLYSTORE_SUBPROCESS_NO_GPU") == "1":
+        class PlaceholderBackend:
+            pass
+        PlaceholderBackend.__name__ = name
+        PlaceholderBackend.__qualname__ = name
+        return PlaceholderBackend
+
+    module_path, class_name = _LAZY_BACKEND_REGISTRY[name]
+    import importlib
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)

src/polystore/async_init.py

@@ -0,0 +1,65 @@
+"""
+Async initialization for I/O backends.
+
+This module provides background loading of GPU-heavy dependencies
+to avoid blocking during import while ensuring they're ready when needed.
+"""
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+_init_complete = threading.Event()
+_init_thread = None
+
+
+def _background_init():
+    """
+    Background thread target for async initialization.
+
+    NOTE: Due to Python's GIL, importing heavy modules in a background thread
+    still blocks the main thread during imports. Therefore, we use lazy initialization
+    for GPU libraries and storage backends - they'll be loaded on first use.
+
+    Currently this is a no-op placeholder. GPU registry initialization happens
+    lazily on first use, not during startup.
+    """
+    try:
+        logger.info("Background I/O initialization complete (lazy mode - no-op)")
+    except Exception as e:
+        logger.error(f"Background I/O initialization failed: {e}")
+    finally:
+        _init_complete.set()
+
+
+def start_async_initialization():
+    """
+    Start background initialization of GPU-heavy dependencies.
+    
+    Call this during application startup (GUI/CLI) to pre-load
+    GPU libraries without blocking. Safe to call multiple times.
+    """
+    global _init_thread
+    
+    if _init_thread is None:
+        _init_thread = threading.Thread(
+            target=_background_init,
+            daemon=True,
+            name="io-async-init"
+        )
+        _init_thread.start()
+        logger.info("Started background I/O initialization")
+
+
+def wait_for_initialization(timeout: float = 60.0) -> bool:
+    """
+    Wait for background initialization to complete.
+    
+    Args:
+        timeout: Maximum seconds to wait
+        
+    Returns:
+        True if completed, False if timeout
+    """
+    return _init_complete.wait(timeout)
+

src/polystore/atomic.py

@@ -122,61 +122,6 @@ def _cleanup_lock(lock_fd: Optional[int], lock_path: Path) -> None:
             logger.warning(f"Error removing lock file {lock_path}: {e}")
 
 
-@contextmanager
-def atomic_write(file_path: Union[str, Path], mode: str = 'w', ensure_directory: bool = True):
-    """
-    Context manager for atomic file writes.
-    
-    Writes to a temporary file first, then renames to the target path.
-    This ensures the operation is atomic - either fully succeeds or fails
-    without leaving partial writes.
-    
-    Args:
-        file_path: Target file path
-        mode: File mode ('w' for text, 'wb' for binary)
-        ensure_directory: Create parent directory if it doesn't exist
-        
-    Example:
-        with atomic_write("output.txt") as f:
-            f.write("data")
-    """
-    file_path = Path(file_path)
-    
-    if ensure_directory:
-        file_path.parent.mkdir(parents=True, exist_ok=True)
-    
-    # Create temporary file in same directory
-    with tempfile.NamedTemporaryFile(
-        mode=mode,
-        dir=file_path.parent,
-        prefix=f"{LOCK_CONFIG.TEMP_PREFIX}{file_path.name}",
-        delete=False
-    ) as tmp_file:
-        tmp_path = tmp_file.name
-        try:
-            yield tmp_file
-            tmp_file.flush()
-            os.fsync(tmp_file.fileno())
-        except Exception:
-            # Clean up temp file on error
-            try:
-                os.unlink(tmp_path)
-            except Exception:
-                pass
-            raise
-    
-    # Atomically replace target file
-    try:
-        os.replace(tmp_path, str(file_path))
-        logger.debug(f"Atomically wrote to {file_path}")
-    except Exception as e:
-        try:
-            os.unlink(tmp_path)
-        except Exception:
-            pass
-        raise FileLockError(f"Atomic write failed for {file_path}: {e}") from e
-
-
 def atomic_write_json(
     file_path: Union[str, Path],
     data: Dict[str, Any],