Runtime regression fix

Author: BrianMichell

src/mdio/segy/_disaster_recovery_wrapper.py

@@ -13,50 +13,6 @@
     from segy.transforms import Transform, TransformPipeline, ByteSwapTransform, IbmFloatTransform
     from numpy.typing import NDArray
 
-
-class HeaderRawTransformedAccessor:
-    """Utility class to access both raw and transformed header data with single filesystem read.
-
-    This class works as a consumer of SegyFile objects without modifying the package.
-    It achieves the goal by:
-    1. Reading raw data from filesystem once
-    2. Applying transforms to get transformed data
-    3. Keeping both versions available
-
-    The transforms used in SEG-Y processing are reversible:
-    - ByteSwapTransform: Self-inverse (swapping twice returns to original)
-    - IbmFloatTransform: Can be reversed by swapping direction
-    """
-
-    def __init__(self, segy_file: SegyFile):
-        """Initialize with a SegyFile instance.
-
-        Args:
-            segy_file: The SegyFile instance to work with
-        """
-        self.segy_file = segy_file
-        self.transform_pipeline = self.segy_file.header.transform_pipeline
-
-    def _reverse_transforms(self, transformed_data: NDArray) -> NDArray:
-        """Reverse the transform pipeline to get raw data from transformed data.
-
-        Args:
-            transformed_data: Data that has been processed through the transform pipeline
-
-        Returns:
-            Raw data equivalent to what was read directly from filesystem
-        """
-        # Start with the transformed data
-        raw_data = transformed_data.copy() if hasattr(transformed_data, 'copy') else transformed_data
-
-
-        # Apply transforms in reverse order with reversed operations
-        for i, transform in enumerate(reversed(self.transform_pipeline.transforms)):
-            raw_data = _reverse_single_transform(raw_data, transform)
-
-        return raw_data
-
-@profile
 def _reverse_single_transform(data: NDArray, transform: Transform) -> NDArray:
     """Reverse a single transform operation.
 
@@ -98,11 +54,10 @@ def _reverse_single_transform(data: NDArray, transform: Transform) -> NDArray:
         # This maintains compatibility if new transforms are added
         return data
 
-
 def get_header_raw_and_transformed(
     segy_file: SegyFile,
     indices: int | list[int] | np.ndarray | slice
-) -> tuple[NDArray, NDArray]:
+) -> tuple[NDArray, NDArray, NDArray]:
     """Convenience function to get both raw and transformed header data.
 
     This is a drop-in replacement that provides the functionality you requested
@@ -127,38 +82,17 @@ def get_header_raw_and_transformed(
         # Slice of headers
         raw_hdrs, transformed_hdrs = get_header_raw_and_transformed(segy_file, slice(0, 10))
     """
-    return _get_header_raw_optimized(segy_file, indices)
-
-@profile
-def _get_header_raw_optimized(
-    segy_file: SegyFile,
-    indices: int | list[int] | np.ndarray | slice
-) -> tuple[NDArray, NDArray]:
-    """Ultra-optimized function that eliminates double disk reads entirely.
 
-    This function:
-    1. Gets transformed headers using the normal API (single disk read)
-    2. Reverses the transforms on the already-loaded data (no second disk read)
-    3. Returns both raw and transformed headers
+    traces = segy_file.trace[indices]
 
-    Args:
-        segy_file: The SegyFile instance
-        indices: Which headers to retrieve
-
-    Returns:
-        Tuple of (raw_headers, transformed_headers) where transformed_headers
-        is the same as what segy_file.header[indices] would return
-    """
-    # Get transformed headers using the normal API (single disk read)
-    transformed_headers = segy_file.header[indices]
+    transformed_headers = traces.header
 
     # Reverse the transforms on the already-loaded transformed data
     # This eliminates the second disk read entirely!
     raw_headers = _reverse_transforms(transformed_headers, segy_file.header.transform_pipeline)
 
-    return raw_headers, transformed_headers
+    return raw_headers, transformed_headers, traces
 
-@profile
 def _reverse_transforms(transformed_data: NDArray, transform_pipeline) -> NDArray:
     """Reverse the transform pipeline to get raw data from transformed data.
 

src/mdio/segy/_workers.py

@@ -121,7 +121,8 @@ def trace_worker(  # noqa: PLR0913
     zarr_config.set({"threading.max_workers": 1})
 
     live_trace_indexes = local_grid_map[not_null].tolist()
-    traces = segy_file.trace[live_trace_indexes]
+    # traces = segy_file.trace[live_trace_indexes]
+    raw_headers, transformed_headers, traces = get_header_raw_and_transformed(segy_file, live_trace_indexes)
 
     header_key = "headers"
     raw_header_key = "raw_headers"
@@ -135,7 +136,7 @@ def trace_worker(  # noqa: PLR0913
         worker_variables.append(raw_header_key)
 
     ds_to_write = dataset[worker_variables]
-    raw_headers, transformed_headers = get_header_raw_and_transformed(segy_file, live_trace_indexes)
+    # raw_headers, transformed_headers = get_header_raw_and_transformed(segy_file, live_trace_indexes)
 
     if header_key in worker_variables:
         # Create temporary array for headers with the correct shape
@@ -153,7 +154,7 @@ def trace_worker(  # noqa: PLR0913
             attrs=ds_to_write[header_key].attrs,
             encoding=ds_to_write[header_key].encoding,  # Not strictly necessary, but safer than not doing it.
         )
-        # del transformed_headers  # Manage memory
+    del transformed_headers  # Manage memory
     if raw_header_key in worker_variables:
         tmp_raw_headers = np.zeros_like(dataset[raw_header_key])
         tmp_raw_headers[not_null] = raw_headers.view("|V240")
@@ -163,8 +164,8 @@ def trace_worker(  # noqa: PLR0913
             attrs=ds_to_write[raw_header_key].attrs,
             encoding=ds_to_write[raw_header_key].encoding,  # Not strictly necessary, but safer than not doing it.
         )
-        del raw_headers  # Manage memory
 
+    del raw_headers  # Manage memory
     data_variable = ds_to_write[data_variable_name]
     fill_value = _get_fill_value(ScalarType(data_variable.dtype.name))
     tmp_samples = np.full_like(data_variable, fill_value=fill_value)

src/mdio/segy/blocked_io.py

@@ -2,7 +2,10 @@
 
 from __future__ import annotations
 
+import multiprocessing as mp
 import os
+from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import as_completed
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -77,20 +80,37 @@ def to_zarr(  # noqa: PLR0913, PLR0915
     chunk_iter = ChunkIterator(shape=data.shape, chunks=worker_chunks, dim_names=data.dims)
     num_chunks = chunk_iter.num_chunks
 
+    # For Unix async writes with s3fs/fsspec & multiprocessing, use 'spawn' instead of default
+    # 'fork' to avoid deadlocks on cloud stores. Slower but necessary. Default on Windows.
+    num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
+    num_workers = min(num_chunks, num_cpus)
+    context = mp.get_context("spawn")
+    executor = ProcessPoolExecutor(max_workers=num_workers, mp_context=context)
+
     segy_kw = {
         "url": segy_file.fs.unstrip_protocol(segy_file.url),
         "spec": segy_file.spec,
         "settings": segy_file.settings,
     }
+    with executor:
+        futures = []
+        common_args = (segy_kw, output_path, data_variable_name)
+        for region in chunk_iter:
+            subset_args = (region, grid_map, dataset.isel(region))
+            future = executor.submit(trace_worker, *common_args, *subset_args)
+            futures.append(future)
+
+        iterable = tqdm(
+            as_completed(futures),
+            total=num_chunks,
+            unit="block",
+            desc="Ingesting traces",
+        )
 
-    common_args = (segy_kw, output_path, data_variable_name)
-
-    # Execute trace_worker serially for profiling
-    for region in tqdm(chunk_iter, total=num_chunks, unit="block", desc="Ingesting traces"):
-        subset_args = (region, grid_map, dataset.isel(region))
-        result = trace_worker(*common_args, *subset_args)
-        if result is not None:
-            _update_stats(final_stats, result)
+        for future in iterable:
+            result = future.result()
+            if result is not None:
+                _update_stats(final_stats, result)
 
     # Xarray doesn't directly support incremental attribute updates when appending to an existing Zarr store.
     # HACK: We will update the array attribute using zarr's API directly.
@@ -260,4 +280,4 @@ def to_segy(
 
         non_consecutive_axes -= 1
 
-    return block_io_records
+    return block_io_records