Strip out local gcs fake server and revert some regressions

Author: BrianMichell

src/mdio/api/io.py

@@ -13,7 +13,6 @@
 from xarray.backends.writers import to_zarr as xr_to_zarr
 
 from mdio.constants import ZarrFormat
-from mdio.core.config import MDIOSettings
 from mdio.core.zarr_io import zarr_warnings_suppress_unstable_structs_v3
 
 if TYPE_CHECKING:
@@ -26,39 +25,13 @@
 
 
 def _normalize_path(path: UPath | Path | str) -> UPath:
-    """Normalize a path to a UPath."""
     return UPath(path)
 
 
 def _normalize_storage_options(path: UPath) -> dict[str, Any] | None:
-    """Normalize storage options from UPath."""
     return None if len(path.storage_options) == 0 else path.storage_options
 
 
-def _get_gcs_store(path: UPath) -> tuple[Any, dict[str, Any] | None]:
-    """Get store and storage options, using local fake GCS server if enabled.
-
-    Args:
-        path: UPath pointing to storage location.
-
-    Returns:
-        Tuple of (store, storage_options) where store is either a mapper or path string.
-    """
-    settings = MDIOSettings()
-
-    if settings.local_gcs_server and str(path).startswith("gs://"):
-        import gcsfs  # noqa: PLC0415
-
-        fs = gcsfs.GCSFileSystem(
-            endpoint_url="http://localhost:4443",
-            token="anon",  # noqa: S106
-        )
-        store = fs.get_mapper(path.as_posix().replace("gs://", ""))
-        return store, None
-
-    return path.as_posix(), _normalize_storage_options(path)
-
-
 def open_mdio(input_path: UPath | Path | str, chunks: T_Chunks = None) -> xr_Dataset:
     """Open a Zarr dataset from the specified universal file path.
 
@@ -114,22 +87,17 @@ def to_mdio(  # noqa: PLR0913
             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")
 
-    store, storage_options = _get_gcs_store(output_path)
-
-    kwargs = {
-        "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)
+        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,
+        )

src/mdio/converters/segy.py

@@ -597,6 +597,7 @@ def segy_to_mdio(  # noqa PLR0913
 
     # Now we can drop them to simplify chunked write of the data variable
     xr_dataset = xr_dataset.drop_vars(drop_vars_delayed)
+
     # 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

src/mdio/core/config.py

@@ -56,10 +56,5 @@ class MDIOSettings(BaseSettings):
         description="Whether to ignore validation checks",
         alias="MDIO_IGNORE_CHECKS",
     )
-    local_gcs_server: bool = Field(
-        default=False,
-        description="Whether to use local fake GCS server for testing (localhost:4443)",
-        alias="MDIO__LOCAL_GCS_SERVER",
-    )
 
     model_config = SettingsConfigDict(case_sensitive=True)

src/mdio/segy/_workers.py

@@ -97,7 +97,7 @@ def trace_worker(  # noqa: PLR0913
     header_array: zarr_Array | None,
     raw_header_array: zarr_Array | None,
     region: dict[str, slice],
-    grid_map_data: np.ndarray,
+    grid_map: zarr_Array,
 ) -> SummaryStatistics | None:
     """Writes a subset of traces from a region of the dataset of Zarr file.
     
@@ -108,7 +108,7 @@ def trace_worker(  # noqa: PLR0913
         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_data: Numpy array mapping live traces to their positions in the dataset.
+        grid_map: Zarr array mapping live traces to their positions in the dataset.
 
     Returns:
         SummaryStatistics object containing statistics about the written traces.
@@ -119,7 +119,7 @@ def trace_worker(  # noqa: PLR0913
     segy_file = _worker_segy_file
 
     region_slices = tuple(region.values())
-    local_grid_map = grid_map_data[region_slices[:-1]]  # minus last (vertical) axis
+    local_grid_map = grid_map[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.

src/mdio/segy/blocked_io.py

@@ -15,7 +15,7 @@
 from tqdm.auto import tqdm
 from zarr import open_group as zarr_open_group
 
-from mdio.api.io import _get_gcs_store
+from mdio.api.io import _normalize_storage_options
 from mdio.builder.schemas.v1.stats import CenteredBinHistogram
 from mdio.builder.schemas.v1.stats import SummaryStatistics
 from mdio.constants import ZarrFormat
@@ -84,10 +84,9 @@ def to_zarr(  # noqa: PLR0913, PLR0915
     zarr_format = zarr.config.get("default_zarr_format")
 
     # Open zarr group once in main process
-    store, storage_options = _get_gcs_store(output_path)
-
+    storage_options = _normalize_storage_options(output_path)
     zarr_group = zarr_open_group(
-        store,
+        output_path.as_posix(),
         mode="r+",
         storage_options=storage_options,
         use_consolidated=zarr_format == ZarrFormat.V2,
@@ -100,10 +99,6 @@ def to_zarr(  # noqa: PLR0913, PLR0915
     header_array = zarr_group.get("headers")
     raw_header_array = zarr_group.get("raw_headers")
 
-    # Convert grid_map to numpy array for serialization
-    # grid_map is an in-memory zarr array, so we can read it all at once
-    grid_map_data = grid_map[:]
-
     # 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_workers = min(num_chunks, settings.import_cpus)
@@ -127,7 +122,7 @@ def to_zarr(  # noqa: PLR0913, PLR0915
                 header_array,
                 raw_header_array,
                 region,
-                grid_map_data,
+                grid_map,
             )
             futures.append(future)