Merge-in latest 'upstream v1'

src/mdio/segy/_workers.py

@@ -4,13 +4,16 @@
 
 import os
 from typing import TYPE_CHECKING
+from typing import TypedDict
 from typing import cast
 
 import numpy as np
 from segy import SegyFile
 
 if TYPE_CHECKING:
     from segy.arrays import HeaderArray
+    from segy.config import SegySettings
+    from segy.schema import SegySpec
     from xarray import Dataset as xr_Dataset
     from zarr import Array as zarr_Array
 
@@ -29,10 +32,7 @@ class SegyFileArguments(TypedDict):
     settings: SegySettings | None
 
 
-def header_scan_worker(
-    segy_kw: SegyFileArguments,
-    trace_range: tuple[int, int],
-) -> HeaderArray:
+def header_scan_worker(segy_kw: SegyFileArguments, trace_range: tuple[int, int]) -> HeaderArray:
     """Header scan worker.
 
     If SegyFile is not open, it can either accept a path string or a handle that was opened in
@@ -69,7 +69,7 @@ def header_scan_worker(
 
 
 def trace_worker_v1(  # noqa: PLR0913
-    segy_file: SegyFile,
+    segy_kw: SegyFileArguments,
     output_location: StorageLocation,
     data_variable_name: str,
     region: dict[str, slice],
@@ -79,7 +79,7 @@ def trace_worker_v1(  # noqa: PLR0913
     """Writes a subset of traces from a region of the dataset of Zarr file.
 
     Args:
-        segy_file: SegyFile instance.
+        segy_kw: Arguments to open SegyFile instance.
         output_location: StorageLocation for the output Zarr dataset
             (e.g. local file path or cloud storage URI) the location
             also includes storage options for cloud storage.
@@ -94,6 +94,10 @@ def trace_worker_v1(  # noqa: PLR0913
     if not dataset.trace_mask.any():
         return None
 
+    # Open the SEG-Y file in every new process / spawned worker since the
+    # open file handles cannot be shared across processes.
+    segy_file = SegyFile(**segy_kw)
+
     not_null = grid_map != UINT32_MAX
 
     live_trace_indexes = grid_map[not_null].tolist()

src/mdio/segy/blocked_io.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import multiprocessing as mp
 import os
 from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import as_completed
@@ -91,15 +92,22 @@ def to_zarr_v1(  # noqa: PLR0913, PLR0915
     # (e.g. for unit testing or for debugging with non-parallelized processing)
     # def _create_executor(num_chunks: int)-> ProcessPoolExecutor:
 
+    # 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)
     # return executor
 
+    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_file, output_location, data_variable_name)
+        common_args = (segy_kw, output_location, data_variable_name)
         for region in chunk_iter:
             index_slices = tuple(region[key] for key in data.dims[:-1])
             subset_args = (