Refactor initial file information handling while avoiding asyncio issues. (#701)

* Fix bug 613 on top of latest main

* Exclude TestBug613HandlesNotClosed from coverage

* Space is added by pre-commit

* Exclude from code coverage: test_teapot_import_cloud_to_cloud

* Remove tests as per the PR review

* PR review: move _get_coordinate_scalar()

* PR Review fixes

* refactor: replace individual SEG-Y file arguments with SegyFileInfo object for cleaner function signatures and consistent data access

* refactor: rename parameter in _get_horizontal_coordinate_unit for consistency with SegyFileInfo usage

* refactor: rename parameter in _add_segy_file_headers for consistency with SegyFileInfo usage

* refactor: rename segy_kw to segy_file_kwargs for consistency with naming conventions

* refactor: rename segy_kw to segy_file_kwargs across functions for improved clarity and consistency

* refactor: reorder imports in scalar.py for adherence to style guidelines and future type checking compatibility

* refactor: simplify info_worker by removing trace_indices handling and unused imports

* refactor: update type hints in _workers.py for improved readability and compliance with TYPE_CHECKING

---------

Co-authored-by: Altay Sansal <[email protected]>

Author: dmitriyrepin

src/mdio/converters/segy.py

@@ -4,13 +4,13 @@
 
 import base64
 import logging
+import multiprocessing as mp
 import os
-from dataclasses import dataclass
+from concurrent.futures import ProcessPoolExecutor
 from typing import TYPE_CHECKING
 
 import numpy as np
 import zarr
-from segy import SegyFile
 from segy.config import SegyFileSettings
 from segy.config import SegyHeaderOverrides
 from segy.standards.codes import MeasurementSystem as SegyMeasurementSystem
@@ -37,9 +37,10 @@
 from mdio.core.utils_write import MAX_SIZE_LIVE_MASK
 from mdio.core.utils_write import get_constrained_chunksize
 from mdio.segy import blocked_io
+from mdio.segy._workers import SegyFileInfo
+from mdio.segy._workers import info_worker
 from mdio.segy.scalar import SCALE_COORDINATE_KEYS
 from mdio.segy.scalar import _apply_coordinate_scalar
-from mdio.segy.scalar import _get_coordinate_scalar
 from mdio.segy.utilities import get_grid_plan
 
 if TYPE_CHECKING:
@@ -54,6 +55,7 @@
     from mdio.builder.schemas import Dataset
     from mdio.builder.templates.abstract_dataset_template import AbstractDatasetTemplate
     from mdio.core.dimension import Dimension
+    from mdio.segy._workers import SegyFileArguments
 
 logger = logging.getLogger(__name__)
 
@@ -135,37 +137,9 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
         raise GridTraceSparsityError(grid.shape, num_traces, msg)
 
 
-@dataclass
-class SegyFileHeaderDump:
-    """Segy metadata information."""
-
-    text_header: str
-    binary_header_dict: dict
-    raw_binary_headers: bytes
-
-
-def _get_segy_file_header_dump(segy_file: SegyFile) -> SegyFileHeaderDump:
-    """Reads information from a SEG-Y file."""
-    text_header = segy_file.text_header
-
-    raw_binary_headers: bytes = segy_file.fs.read_block(
-        fn=segy_file.url,
-        offset=segy_file.spec.binary_header.offset,
-        length=segy_file.spec.binary_header.itemsize,
-    )
-
-    # We read here twice, but it's ok for now. Only 400-bytes.
-    binary_header_dict = segy_file.binary_header.to_dict()
-
-    return SegyFileHeaderDump(
-        text_header=text_header,
-        binary_header_dict=binary_header_dict,
-        raw_binary_headers=raw_binary_headers,
-    )
-
-
 def _scan_for_headers(
-    segy_file: SegyFile,
+    segy_file_kwargs: SegyFileArguments,
+    segy_file_info: SegyFileInfo,
     template: AbstractDatasetTemplate,
     grid_overrides: dict[str, Any] | None = None,
 ) -> tuple[list[Dimension], SegyHeaderArray]:
@@ -176,7 +150,8 @@ def _scan_for_headers(
     """
     full_chunk_size = template.full_chunk_size
     segy_dimensions, chunk_size, segy_headers = get_grid_plan(
-        segy_file=segy_file,
+        segy_file_kwargs=segy_file_kwargs,
+        segy_file_info=segy_file_info,
         return_headers=True,
         template=template,
         chunksize=full_chunk_size,
@@ -192,12 +167,29 @@ def _scan_for_headers(
     return segy_dimensions, segy_headers
 
 
-def _build_and_check_grid(segy_dimensions: list[Dimension], num_traces: int, segy_headers: SegyHeaderArray) -> Grid:
+def _read_segy_file_info(segy_file_kwargs: SegyFileArguments) -> SegyFileInfo:
+    """Read SEG-Y file in a separate process.
+
+    This is an ugly workaround for Zarr issues 3487 'Explicitly using fsspec and zarr FsspecStore causes
+    RuntimeError "Task attached to a different loop"'
+    """
+    # TODO (Dmitriy Repin): when Zarr issue 3487 is resolved, we can remove this workaround
+    # https://github.com/zarr-developers/zarr-python/issues/3487
+    with ProcessPoolExecutor(max_workers=1, mp_context=mp.get_context("spawn")) as executor:
+        future = executor.submit(info_worker, segy_file_kwargs)
+        return future.result()
+
+
+def _build_and_check_grid(
+    segy_dimensions: list[Dimension],
+    segy_file_info: SegyFileInfo,
+    segy_headers: SegyHeaderArray,
+) -> Grid:
     """Build and check the grid from the SEG-Y headers and dimensions.
 
     Args:
         segy_dimensions: List of of all SEG-Y dimensions to build grid from.
-        num_traces: Number of traces in the SEG-Y file.
+        segy_file_info: SegyFileInfo instance containing the SEG-Y file information.
         segy_headers: Headers read in from SEG-Y file for building the trace map.
 
     Returns:
@@ -207,6 +199,7 @@ def _build_and_check_grid(segy_dimensions: list[Dimension], num_traces: int, seg
         GridTraceCountError: If number of traces in SEG-Y file does not match the parsed grid
     """
     grid = Grid(dims=segy_dimensions)
+    num_traces = segy_file_info.num_traces
     grid_density_qc(grid, num_traces)
     grid.build_map(segy_headers)
     # Check grid validity by comparing trace numbers
@@ -303,9 +296,9 @@ def populate_non_dim_coordinates(
     return dataset, drop_vars_delayed
 
 
-def _get_horizontal_coordinate_unit(segy_info: SegyFileHeaderDump) -> LengthUnitModel | None:
+def _get_horizontal_coordinate_unit(segy_file_info: SegyFileInfo) -> LengthUnitModel | None:
     """Get the coordinate unit from the SEG-Y headers."""
-    measurement_system_code = int(segy_info.binary_header_dict[MEASUREMENT_SYSTEM_KEY])
+    measurement_system_code = int(segy_file_info.binary_header_dict[MEASUREMENT_SYSTEM_KEY])
 
     if measurement_system_code not in (1, 2):
         logger.warning(
@@ -359,19 +352,19 @@ def _populate_coordinates(
     return dataset, drop_vars_delayed
 
 
-def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_header_dump: SegyFileHeaderDump) -> xr_Dataset:
+def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo) -> xr_Dataset:
     save_file_header = os.getenv("MDIO__IMPORT__SAVE_SEGY_FILE_HEADER", "") in ("1", "true", "yes", "on")
     if not save_file_header:
         return xr_dataset
 
     expected_rows = 40
     expected_cols = 80
 
-    text_header_rows = segy_file_header_dump.text_header.splitlines()
+    text_header_rows = segy_file_info.text_header.splitlines()
     text_header_cols_bad = [len(row) != expected_cols for row in text_header_rows]
 
     if len(text_header_rows) != expected_rows:
-        err = f"Invalid text header count: expected {expected_rows}, got {len(segy_file_header_dump.text_header)}"
+        err = f"Invalid text header count: expected {expected_rows}, got {len(segy_file_info.text_header)}"
         raise ValueError(err)
 
     if any(text_header_cols_bad):
@@ -381,12 +374,12 @@ def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_header_dump: SegyFi
     xr_dataset["segy_file_header"] = ((), "")
     xr_dataset["segy_file_header"].attrs.update(
         {
-            "textHeader": segy_file_header_dump.text_header,
-            "binaryHeader": segy_file_header_dump.binary_header_dict,
+            "textHeader": segy_file_info.text_header,
+            "binaryHeader": segy_file_info.binary_header_dict,
         }
     )
     if os.getenv("MDIO__IMPORT__RAW_HEADERS") in ("1", "true", "yes", "on"):
-        raw_binary_base64 = base64.b64encode(segy_file_header_dump.raw_binary_headers).decode("ascii")
+        raw_binary_base64 = base64.b64encode(segy_file_info.raw_binary_headers).decode("ascii")
         xr_dataset["segy_file_header"].attrs.update({"rawBinaryHeader": raw_binary_base64})
 
     return xr_dataset
@@ -532,17 +525,21 @@ def segy_to_mdio(  # noqa PLR0913
         raise FileExistsError(err)
 
     segy_settings = SegyFileSettings(storage_options=input_path.storage_options)
-    segy_file = SegyFile(
-        url=input_path.as_posix(),
-        spec=segy_spec,
-        settings=segy_settings,
-        header_overrides=segy_header_overrides,
+    segy_file_kwargs: SegyFileArguments = {
+        "url": input_path.as_posix(),
+        "spec": segy_spec,
+        "settings": segy_settings,
+        "header_overrides": segy_header_overrides,
+    }
+    segy_file_info = _read_segy_file_info(segy_file_kwargs)
+
+    segy_dimensions, segy_headers = _scan_for_headers(
+        segy_file_kwargs,
+        segy_file_info,
+        template=mdio_template,
+        grid_overrides=grid_overrides,
     )
-    segy_info: SegyFileHeaderDump = _get_segy_file_header_dump(segy_file)
-
-    segy_dimensions, segy_headers = _scan_for_headers(segy_file, mdio_template, grid_overrides)
-
-    grid = _build_and_check_grid(segy_dimensions, segy_file.num_traces, segy_headers)
+    grid = _build_and_check_grid(segy_dimensions, segy_file_info, segy_headers)
 
     _, non_dim_coords = _get_coordinates(grid, segy_headers, mdio_template)
     header_dtype = to_structured_type(segy_spec.trace.header.dtype)
@@ -554,7 +551,7 @@ def segy_to_mdio(  # noqa PLR0913
             logger.warning("MDIO__IMPORT__RAW_HEADERS is experimental and expected to change or be removed.")
             mdio_template = _add_raw_headers_to_template(mdio_template)
 
-    horizontal_unit = _get_horizontal_coordinate_unit(segy_info)
+    horizontal_unit = _get_horizontal_coordinate_unit(segy_file_info)
     mdio_ds: Dataset = mdio_template.build_dataset(
         name=mdio_template.name,
         sizes=grid.shape,
@@ -571,15 +568,14 @@ def segy_to_mdio(  # noqa PLR0913
 
     xr_dataset: xr_Dataset = to_xarray_dataset(mdio_ds=mdio_ds)
 
-    coordinate_scalar = _get_coordinate_scalar(segy_file)
     xr_dataset, drop_vars_delayed = _populate_coordinates(
         dataset=xr_dataset,
         grid=grid,
         coords=non_dim_coords,
-        horizontal_coordinate_scalar=coordinate_scalar,
+        horizontal_coordinate_scalar=segy_file_info.coordinate_scalar,
     )
 
-    xr_dataset = _add_segy_file_headers(xr_dataset, segy_info)
+    xr_dataset = _add_segy_file_headers(xr_dataset, segy_file_info)
 
     xr_dataset.trace_mask.data[:] = grid.live_mask
     # IMPORTANT: Do not drop the "trace_mask" here, as it will be used later in
@@ -600,7 +596,7 @@ def segy_to_mdio(  # noqa PLR0913
     # This is an memory-expensive and time-consuming read-write operation
     # performed in chunks to save the memory
     blocked_io.to_zarr(
-        segy_file=segy_file,
+        segy_file_kwargs=segy_file_kwargs,
         output_path=output_path,
         grid_map=grid.map,
         dataset=xr_dataset,

src/mdio/segy/_workers.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+import logging
 import os
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 from typing import TypedDict
 
@@ -13,9 +15,11 @@
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.dtype import ScalarType
 from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
+from mdio.segy.scalar import _get_coordinate_scalar
 
 if TYPE_CHECKING:
     from segy.config import SegyFileSettings
+    from segy.config import SegyHeaderOverrides
     from segy.schema import SegySpec
     from upath import UPath
     from xarray import Dataset as xr_Dataset
@@ -29,32 +33,41 @@
 from mdio.builder.xarray_builder import _get_fill_value
 from mdio.constants import fill_value_map
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
+logger = logging.getLogger(__name__)
+
 
 class SegyFileArguments(TypedDict):
     """Arguments to open SegyFile instance creation."""
 
     url: str
     spec: SegySpec | None
     settings: SegyFileSettings | None
+    header_overrides: SegyHeaderOverrides | None
 
 
 def header_scan_worker(
-    segy_kw: SegyFileArguments, trace_range: tuple[int, int], subset: list[str] | None = None
+    segy_file_kwargs: SegyFileArguments,
+    trace_range: tuple[int, int],
+    subset: list[str] | None = None,
 ) -> HeaderArray:
     """Header scan worker.
 
     If SegyFile is not open, it can either accept a path string or a handle that was opened in
     a different context manager.
 
     Args:
-        segy_kw: Arguments to open SegyFile instance.
+        segy_file_kwargs: Arguments to open SegyFile instance.
         trace_range: Tuple consisting of the trace ranges to read.
         subset: List of header names to filter and keep.
 
     Returns:
         HeaderArray parsed from SEG-Y library.
     """
-    segy_file = SegyFile(**segy_kw)
+    segy_file = SegyFile(**segy_file_kwargs)
 
     slice_ = slice(*trace_range)
 
@@ -82,7 +95,7 @@ def header_scan_worker(
 
 
 def trace_worker(  # noqa: PLR0913
-    segy_kw: SegyFileArguments,
+    segy_file_kwargs: SegyFileArguments,
     output_path: UPath,
     data_variable_name: str,
     region: dict[str, slice],
@@ -92,7 +105,7 @@ def trace_worker(  # noqa: PLR0913
     """Writes a subset of traces from a region of the dataset of Zarr file.
 
     Args:
-        segy_kw: Arguments to open SegyFile instance.
+        segy_file_kwargs: Arguments to open SegyFile instance.
         output_path: Universal Path for the output Zarr dataset
             (e.g. local file path or cloud storage URI) the location
             also includes storage options for cloud storage.
@@ -114,7 +127,7 @@ def trace_worker(  # noqa: PLR0913
         return None
 
     # Open the SEG-Y file in this process since the open file handles cannot be shared across processes.
-    segy_file = SegyFile(**segy_kw)
+    segy_file = SegyFile(**segy_file_kwargs)
 
     # Setting the zarr config to 1 thread to ensure we honor the `MDIO__IMPORT__MAX_WORKERS` environment variable.
     # The Zarr 3 engine utilizes multiple threads. This can lead to resource contention and unpredictable memory usage.
@@ -196,3 +209,52 @@ def trace_worker(  # noqa: PLR0913
         sum_squares=(np.ma.power(nonzero_samples, 2).sum(dtype="float64")),
         histogram=histogram,
     )
+
+
+@dataclass
+class SegyFileInfo:
+    """SEG-Y file header information."""
+
+    num_traces: int
+    sample_labels: NDArray[np.int32]
+    text_header: str
+    binary_header_dict: dict
+    raw_binary_headers: bytes
+    coordinate_scalar: int
+
+
+def info_worker(segy_file_kwargs: SegyFileArguments) -> SegyFileInfo:
+    """Reads information from a SEG-Y file.
+
+    Args:
+        segy_file_kwargs: Arguments to open SegyFile instance.
+
+    Returns:
+        SegyFileInfo containing number of traces, sample labels, and header info.
+    """
+    segy_file = SegyFile(**segy_file_kwargs)
+    num_traces = segy_file.num_traces
+    sample_labels = segy_file.sample_labels
+
+    text_header = segy_file.text_header
+
+    # Get header information directly
+    raw_binary_headers = segy_file.fs.read_block(
+        fn=segy_file.url,
+        offset=segy_file.spec.binary_header.offset,
+        length=segy_file.spec.binary_header.itemsize,
+    )
+
+    # We read here twice, but it's ok for now. Only 400-bytes.
+    binary_header_dict = segy_file.binary_header.to_dict()
+
+    coordinate_scalar = _get_coordinate_scalar(segy_file)
+
+    return SegyFileInfo(
+        num_traces=num_traces,
+        sample_labels=sample_labels,
+        text_header=text_header,
+        binary_header_dict=binary_header_dict,
+        raw_binary_headers=raw_binary_headers,
+        coordinate_scalar=coordinate_scalar,
+    )