Working tests with fake server

BrianMichell · BrianMichell · commit 515ca6c5f49d · 2025-10-31T18:39:40.000Z
diff --git a/src/mdio/api/io.py b/src/mdio/api/io.py
@@ -23,13 +23,46 @@
     from xarray.core.types import T_Chunks
     from xarray.core.types import ZarrWriteModes
 
-
 def _normalize_path(path: UPath | Path | str) -> UPath:
-    return UPath(path)
+    """Normalize a path to a UPath.
 
+    For gs:// paths, the fake GCS server configuration is handled via storage_options
+    in _normalize_storage_options().
+    """
+    from upath import UPath
+
+    return UPath(path)
 
 def _normalize_storage_options(path: UPath) -> dict[str, Any] | None:
-    return None if len(path.storage_options) == 0 else path.storage_options
+    """Normalize and patch storage options for UPath paths.
+
+    - Extracts any existing options from the UPath.
+    - Automatically redirects gs:// URLs to a local fake-GCS endpoint
+      when testing (localhost:4443).
+    """
+    import gcsfs
+
+    # Start with any existing options from UPath
+    storage_options = dict(path.storage_options) if len(path.storage_options) else {}
+
+    # Redirect gs:// to local fake-GCS server for testing
+    if str(path).startswith("gs://"):
+        fs = gcsfs.GCSFileSystem(
+            endpoint_url="http://localhost:4443",
+            token="anon",
+        )
+        base_url = getattr(getattr(fs, "session", None), "_base_url", "http://localhost:4443")
+        print(f"[mdio.utils] Redirecting GCS path to local fake server: {base_url}")
+        storage_options["fs"] = fs
+
+    return storage_options or None
+
+# def _normalize_path(path: UPath | Path | str) -> UPath:
+#     return UPath(path)
+
+
+# def _normalize_storage_options(path: UPath) -> dict[str, Any] | None:
+#     return None if len(path.storage_options) == 0 else path.storage_options
 
 
 def open_mdio(input_path: UPath | Path | str, chunks: T_Chunks = None) -> xr_Dataset:
@@ -49,6 +82,8 @@ def open_mdio(input_path: UPath | Path | str, chunks: T_Chunks = None) -> xr_Dat
     Returns:
         An Xarray dataset opened from the input path.
     """
+    import zarr
+
     input_path = _normalize_path(input_path)
     storage_options = _normalize_storage_options(input_path)
     zarr_format = zarr.config.get("default_zarr_format")
@@ -61,43 +96,101 @@ def open_mdio(input_path: UPath | Path | str, chunks: T_Chunks = None) -> xr_Dat
         consolidated=zarr_format == ZarrFormat.V2,  # on for v2, off for v3
     )
 
-
-def to_mdio(  # noqa: PLR0913
+def to_mdio(
     dataset: Dataset,
     output_path: UPath | Path | str,
     mode: ZarrWriteModes | None = None,
     *,
     compute: bool = True,
-    region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None = None,
-) -> None:
-    """Write dataset contents to an MDIO output_path.
+    region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None = None,):
+    """Write dataset contents to an MDIO output_path."""
+    import gcsfs
+    import zarr
 
-    Args:
-        dataset: The dataset to write.
-        output_path: The universal path of the output MDIO file.
-        mode: Persistence mode: "w" means create (overwrite if exists)
-            "w-" means create (fail if exists)
-            "a" means override all existing variables including dimension coordinates (create if does not exist)
-            "a-" means only append those variables that have ``append_dim``.
-            "r+" means modify existing array *values* only (raise an error if any metadata or shapes would change).
-            The default mode is "r+" if ``region`` is set and ``w-`` otherwise.
-        compute: If True write array data immediately; otherwise return a ``dask.delayed.Delayed`` object that
-            can be computed to write array data later. Metadata is always updated eagerly.
-        region: Optional mapping from dimension names to either a) ``"auto"``, or b) integer slices, indicating
-            the region of existing MDIO array(s) in which to write this dataset's data.
-    """
     output_path = _normalize_path(output_path)
-    storage_options = _normalize_storage_options(output_path)
     zarr_format = zarr.config.get("default_zarr_format")
 
-    with zarr_warnings_suppress_unstable_structs_v3():
-        xr_to_zarr(
-            dataset,
-            store=output_path.as_posix(),  # xarray doesn't like URI when file:// is protocol
-            mode=mode,
-            compute=compute,
-            consolidated=zarr_format == ZarrFormat.V2,  # on for v2, off for v3
-            region=region,
-            storage_options=storage_options,
-            write_empty_chunks=False,
+    # For GCS paths, create FSMap for fake GCS server
+    if str(output_path).startswith("gs://"):
+        fs = gcsfs.GCSFileSystem(
+            endpoint_url="http://localhost:4443",
+            token="anon",
         )
+        base_url = getattr(getattr(fs, "session", None), "_base_url", "http://localhost:4443")
+        print(f"[mdio.utils] Using fake GCS mapper via {base_url}")
+        store = fs.get_mapper(output_path.as_posix().replace("gs://", ""))
+        storage_options = None  # Must be None when passing a mapper
+    else:
+        store = output_path.as_posix()
+        storage_options = _normalize_storage_options(output_path)
+
+    print(f"[mdio.utils] Writing to store: {store}")
+    print(f"[mdio.utils] Storage options: {storage_options}")
+
+    kwargs = dict(
+        dataset=dataset,
+        store=store,
+        mode=mode,
+        compute=compute,
+        consolidated=zarr_format == ZarrFormat.V2,
+        region=region,
+        write_empty_chunks=False,
+    )
+    if storage_options is not None and not isinstance(store, dict):
+        kwargs["storage_options"] = storage_options
+
+    with zarr_warnings_suppress_unstable_structs_v3():
+        xr_to_zarr(**kwargs)
+
+
+# def to_mdio(  # noqa: PLR0913
+#     dataset: Dataset,
+#     output_path: UPath | Path | str,
+#     mode: ZarrWriteModes | None = None,
+#     *,
+#     compute: bool = True,
+#     region: Mapping[str, slice | Literal["auto"]] | Literal["auto"] | None = None,
+# ) -> None:
+#     """Write dataset contents to an MDIO output_path."""
+#     import gcsfs, zarr
+
+#     output_path = _normalize_path(output_path)
+
+#     if output_path.as_posix().startswith("gs://"):
+#         fs = gcsfs.GCSFileSystem(
+#             endpoint_url="http://localhost:4443",
+#             token="anon",
+#         )
+
+#         base_url = getattr(getattr(fs, "session", None), "_base_url", "http://localhost:4443")
+#         print(f"Using custom fake GCS filesystem with endpoint {base_url}")
+
+#         # Build a mapper so all I/O uses the fake GCS server
+#         mapper = fs.get_mapper(output_path.as_posix().replace("gs://", ""))
+#         store = mapper
+#         storage_options = None  # must be None when passing a mapper
+#     else:
+#         store = output_path.as_posix()
+#         storage_options = _normalize_storage_options(output_path) or {}
+
+#     print(f"Writing to store: {store}")
+#     zarr_format = zarr.config.get("default_zarr_format")
+
+#     kwargs = dict(
+#         dataset=dataset,
+#         store=store,
+#         mode=mode,
+#         compute=compute,
+#         consolidated=zarr_format == ZarrFormat.V2,
+#         region=region,
+#         write_empty_chunks=False,
+#     )
+#     if storage_options is not None:
+#         kwargs["storage_options"] = storage_options
+
+#     with zarr_warnings_suppress_unstable_structs_v3():
+#         xr_to_zarr(**kwargs)
+
+
+
+
diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py
@@ -531,9 +531,11 @@ def segy_to_mdio(  # noqa PLR0913
     input_path = _normalize_path(input_path)
     output_path = _normalize_path(output_path)
 
+    print("Checking if output path exists...")
     if not overwrite and output_path.exists():
         err = f"Output location '{output_path.as_posix()}' exists. Set `overwrite=True` if intended."
         raise FileExistsError(err)
+    print("Output path check was fine")
 
     segy_settings = SegyFileSettings(storage_options=input_path.storage_options)
     segy_file_kwargs: SegyFileArguments = {
@@ -589,15 +591,16 @@ def segy_to_mdio(  # noqa PLR0913
     # blocked_io.to_zarr() -> _workers.trace_worker()
 
     # This will create the Zarr store with the correct structure but with empty arrays
+    print("Creating Zarr store...")
     to_mdio(xr_dataset, output_path=output_path, mode="w", compute=False)
-
+    print("Zarr store created")
     # This will write the non-dimension coordinates and trace mask
     meta_ds = xr_dataset[drop_vars_delayed + ["trace_mask"]]
     to_mdio(meta_ds, output_path=output_path, mode="r+", compute=True)
-
+    print("Non-dimension coordinates and trace mask written")
     # Now we can drop them to simplify chunked write of the data variable
     xr_dataset = xr_dataset.drop_vars(drop_vars_delayed)
-
+    print("Dropped variables")
     # Write the headers and traces in chunks using grid_map to indicate dead traces
     default_variable_name = mdio_template.default_variable_name
     # This is an memory-expensive and time-consuming read-write operation
diff --git a/src/mdio/segy/_workers.py b/src/mdio/segy/_workers.py
@@ -8,25 +8,47 @@
 import numpy as np
 from segy.arrays import HeaderArray
 
-from mdio.api.io import _normalize_storage_options
 from mdio.core.config import MDIOSettings
 from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
 from mdio.segy.file import SegyFileArguments
 from mdio.segy.file import SegyFileWrapper
 
 if TYPE_CHECKING:
-    from upath import UPath
     from zarr import Array as zarr_Array
+    from zarr import Group as zarr_Group
+    from segy import SegyFile
 
-from zarr import open_group as zarr_open_group
 from zarr.core.config import config as zarr_config
-
 from mdio.builder.schemas.v1.stats import CenteredBinHistogram
 from mdio.builder.schemas.v1.stats import SummaryStatistics
 from mdio.constants import fill_value_map
 
 logger = logging.getLogger(__name__)
 
+# Global variable to store opened segy file per worker process
+_worker_segy_file = None
+
+
+def _init_worker(segy_file_kwargs: SegyFileArguments) -> None:
+    """Initialize worker process with persistent segy file handle.
+    
+    This function is called once per worker process to open the segy file,
+    which is then reused across all tasks in that worker.
+    
+    Args:
+        segy_file_kwargs: Arguments to open SegyFile instance.
+    """
+    global _worker_segy_file
+    
+    from segy import SegyFile
+
+    # 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.
+    zarr_config.set({"threading.max_workers": 1})
+
+    # Open the SEG-Y file once per worker
+    _worker_segy_file = SegyFile(**segy_file_kwargs)
+
 
 def header_scan_worker(
     segy_file_kwargs: SegyFileArguments,
@@ -71,55 +93,42 @@ def header_scan_worker(
 
 
 def trace_worker(  # noqa: PLR0913
-    segy_file_kwargs: SegyFileArguments,
-    output_path: UPath,
-    data_variable_name: str,
+    data_array: zarr_Array,
+    header_array: zarr_Array | None,
+    raw_header_array: zarr_Array | None,
     region: dict[str, slice],
-    grid_map: zarr_Array,
+    grid_map_data: np.ndarray,
 ) -> SummaryStatistics | None:
     """Writes a subset of traces from a region of the dataset of Zarr file.
+    
+    Uses pre-opened segy file from _init_worker and receives zarr arrays directly.
 
     Args:
-        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.
-        data_variable_name: Name of the data variable to write.
+        data_array: Zarr array for writing trace data.
+        header_array: Zarr array for writing trace headers (or None if not needed).
+        raw_header_array: Zarr array for writing raw headers (or None if not needed).
         region: Region of the dataset to write to.
-        grid_map: Zarr array mapping live traces to their positions in the dataset.
+        grid_map_data: Numpy array mapping live traces to their positions in the dataset.
 
     Returns:
         SummaryStatistics object containing statistics about the written traces.
     """
+    global _worker_segy_file
+    
+    # Use the pre-opened segy file from worker initialization
+    segy_file = _worker_segy_file
+
     region_slices = tuple(region.values())
-    local_grid_map = grid_map[region_slices[:-1]]  # minus last (vertical) axis
+    local_grid_map = grid_map_data[region_slices[:-1]]  # minus last (vertical) axis
 
     # The dtype.max is the sentinel value for the grid map.
     # Normally, this is uint32, but some grids need to be promoted to uint64.
     not_null = local_grid_map != fill_value_map.get(local_grid_map.dtype.name)
     if not not_null.any():
         return None
 
-    # Open the SEG-Y file in this process since the open file handles cannot be shared across processes.
-    segy_file = SegyFileWrapper(**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.
-    zarr_config.set({"threading.max_workers": 1})
-
     live_trace_indexes = local_grid_map[not_null].tolist()
 
-    # Open the zarr group to write directly
-    storage_options = _normalize_storage_options(output_path)
-    zarr_group = zarr_open_group(output_path.as_posix(), mode="r+", storage_options=storage_options)
-
-    header_key = "headers"
-    raw_header_key = "raw_headers"
-
-    # Check which variables exist in the zarr store
-    available_arrays = list(zarr_group.array_keys())
-
-    # traces = segy_file.trace[live_trace_indexes]
     # Raw headers are not intended to remain as a feature of the SEGY ingestion.
     # For that reason, we have wrapped the accessors to provide an interface that can be removed
     # and not require additional changes to the below code.
@@ -132,24 +141,21 @@ def trace_worker(  # noqa: PLR0913
     full_shape = tuple(s.stop - s.start for s in region_slices)
     header_shape = tuple(s.stop - s.start for s in header_region_slices)
 
-    # Write raw headers if they exist
+    # Write raw headers if array was provided
     # Headers only have spatial dimensions (no sample dimension)
-    if raw_header_key in available_arrays:
-        raw_header_array = zarr_group[raw_header_key]
+    if raw_header_array is not None:
         tmp_raw_headers = np.full(header_shape, raw_header_array.fill_value)
         tmp_raw_headers[not_null] = traces.raw_header
         raw_header_array[header_region_slices] = tmp_raw_headers
 
-    # Write headers if they exist
+    # Write headers if array was provided
     # Headers only have spatial dimensions (no sample dimension)
-    if header_key in available_arrays:
-        header_array = zarr_group[header_key]
+    if header_array is not None:
         tmp_headers = np.full(header_shape, header_array.fill_value)
         tmp_headers[not_null] = traces.header
         header_array[header_region_slices] = tmp_headers
 
     # Write the data variable
-    data_array = zarr_group[data_variable_name]
     tmp_samples = np.full(full_shape, data_array.fill_value)
     tmp_samples[not_null] = traces.sample
     data_array[region_slices] = tmp_samples
diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py
