From d27879f5fdea28aa07b45094a9dc4b512e847a62 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Mon, 2 Jun 2025 16:15:29 +0000 Subject: [PATCH 01/16] Begin testing for lazy compute of grid --- src/mdio/converters/segy.py | 76 ++++++++++++++++++---- src/mdio/core/grid.py | 96 ++++++++++++++++----------- src/mdio/segy/_workers.py | 66 ++++++++++--------- src/mdio/segy/blocked_io.py | 125 ++++++++++++++++++++++++++---------- 4 files changed, 251 insertions(+), 112 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 15946e84c..7e821fc76 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -354,6 +354,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ... grid_overrides={"HasDuplicates": True}, ... ) """ + print("Entering segy_to_mdio") index_names = index_names or [f"dim_{i}" for i in range(len(index_bytes))] index_types = index_types or ["int32"] * len(index_bytes) @@ -368,6 +369,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 storage_options_input = storage_options_input or {} storage_options_output = storage_options_output or {} + print("pre-setup") # Open SEG-Y with MDIO's SegySpec. Endianness will be inferred. mdio_spec = mdio_segy_spec() segy_settings = SegySettings(storage_options=storage_options_input) @@ -377,13 +379,15 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 binary_header = segy.binary_header num_traces = segy.num_traces + print("pre-index") # Index the dataset using a spec that interprets the user provided index headers. - index_fields = [] + index_fields: list[HeaderField] = [] for name, byte, format_ in zip(index_names, index_bytes, index_types, strict=True): index_fields.append(HeaderField(name=name, byte=byte, format=format_)) mdio_spec_grid = mdio_spec.customize(trace_header_fields=index_fields) segy_grid = SegyFile(url=segy_path, spec=mdio_spec_grid, settings=segy_settings) + print("pre-get_grid_plan") dimensions, chunksize, index_headers = get_grid_plan( segy_file=segy_grid, return_headers=True, @@ -391,17 +395,27 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 grid_overrides=grid_overrides, ) grid = Grid(dims=dimensions) + print("pre-grid_density_qc") grid_density_qc(grid, num_traces) + print("pre-build_map") grid.build_map(index_headers) - # Check grid validity by comparing trace numbers - if np.sum(grid.live_mask) != num_traces: + print("pre-valid_mask") + # Check grid validity by ensuring every trace's header-index is within dimension bounds + valid_mask = np.ones(grid.num_traces, dtype=bool) + for d_idx in range(len(grid.header_index_arrays)): + coords = grid.header_index_arrays[d_idx] + valid_mask &= (coords < grid.shape[d_idx]) + valid_count = int(np.count_nonzero(valid_mask)) + if valid_count != num_traces: for dim_name in grid.dim_names: - dim_min, dim_max = grid.get_min(dim_name), grid.get_max(dim_name) + dim_min = grid.get_min(dim_name) + dim_max = grid.get_max(dim_name) logger.warning("%s min: %s max: %s", dim_name, dim_min, dim_max) logger.warning("Ingestion grid shape: %s.", grid.shape) - raise GridTraceCountError(np.sum(grid.live_mask), num_traces) + raise GridTraceCountError(valid_count, num_traces) + print("pre-chunksize") if chunksize is None: dim_count = len(index_names) + 1 if dim_count == 2: # noqa: PLR2004 @@ -424,6 +438,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None] suffix = "".join(suffix) + print("pre-compressors") compressors = get_compressor(lossless, compression_tolerance) header_dtype = segy.spec.trace.header.dtype.newbyteorder("=") var_conf = MDIOVariableConfig( @@ -435,6 +450,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ) config = MDIOCreateConfig(path=mdio_path_or_buffer, grid=grid, variables=[var_conf]) + print("pre-create_empty") root_group = create_empty( config, overwrite=overwrite, @@ -446,23 +462,61 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 data_array = data_group[f"chunked_{suffix}"] header_array = meta_group[f"chunked_{suffix}_trace_headers"] - # Write actual live mask and metadata to empty MDIO - meta_group["live_mask"][:] = grid.live_mask[:] - nonzero_count = np.count_nonzero(grid.live_mask) + print("pre-live_mask") + live_mask_array = meta_group["live_mask"] + # 'live_mask_array' has the same first N–1 dims as 'grid.shape[:-1]' + # Build a ChunkIterator over the live_mask (no sample axis) + from mdio.core.indexing import ChunkIterator + + chunker = ChunkIterator(live_mask_array, chunk_samples=False) + for chunk_indices in chunker: + # chunk_indices is a tuple of N–1 slice objects + trace_ids = grid.get_traces_for_chunk(chunk_indices) + if trace_ids.size == 0: + continue + + # Build a temporary boolean block of shape = chunk shape + block_shape = tuple(sl.stop - sl.start for sl in chunk_indices) + block = np.zeros(block_shape, dtype=bool) + + # Compute local coords within this block for each trace_id + local_coords: list[np.ndarray] = [] + for dim_idx, sl in enumerate(chunk_indices): + hdr_arr = grid.header_index_arrays[dim_idx] + local_idx = (hdr_arr[trace_ids] - sl.start).astype(int) + local_coords.append(local_idx) + + # Mark live cells in the temporary block + block[tuple(local_coords)] = True + + # Write the entire block to Zarr at once + live_mask_array.set_basic_selection(selection=chunk_indices, value=block) + + nonzero_count = grid.num_traces + + print("pre-write_attribute") write_attribute(name="trace_count", zarr_group=root_group, attribute=nonzero_count) write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n")) write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) + print("pre-to_zarr") # Write traces + zarr_root = mdio_path_or_buffer # the same path you passed earlier to create_empty + data_var = f"data/chunked_{suffix}" + header_var = f"metadata/chunked_{suffix}_trace_headers" + stats = blocked_io.to_zarr( segy_file=segy, grid=grid, - data_array=data_array, - header_array=header_array, + zarr_root_path=zarr_root, + data_var_path=data_var, + header_var_path=header_var, ) + print("pre-write_attribute") # Write actual stats for key, value in stats.items(): write_attribute(name=key, zarr_group=root_group, attribute=value) - zarr.consolidate_metadata(root_group.store) + print("pre-consolidate_metadata") + zarr.consolidate_metadata(root_group.store) \ No newline at end of file diff --git a/src/mdio/core/grid.py b/src/mdio/core/grid.py index a81d0cd03..9ed9378c1 100644 --- a/src/mdio/core/grid.py +++ b/src/mdio/core/grid.py @@ -65,6 +65,9 @@ def __post_init__(self) -> None: self.dim_names = tuple(dim.name for dim in self.dims) self.shape = tuple(dim.size for dim in self.dims) self.ndim = len(self.dims) + # Prepare attributes for lazy mapping; they will be set in build_map + self.header_index_arrays: tuple[np.ndarray, ...] = () + self.num_traces: int = 0 def __getitem__(self, item: int) -> Dimension: """Get a dimension by index.""" @@ -106,47 +109,64 @@ def from_zarr(cls, zarr_root: zarr.Group) -> Grid: return cls(dims_list) def build_map(self, index_headers: HeaderArray) -> None: - """Build trace mapping and live mask from header indices. + """Compute per-trace grid coordinates (lazy map). + + Instead of allocating a full `self.map` and `self.live_mask`, this computes, for each trace, + its integer index along each dimension (excluding the final sample dimension) and stores them in + `self.header_index_arrays`. The full mapping can then be derived chunk-by-chunk when writing. Args: - index_headers: Header array containing dimension indices. + index_headers: Header array containing dimension indices (length = number of traces). + """ + # Number of traces in the SEG-Y + self.num_traces = int(index_headers.shape[0]) + + # For each dimension except the final sample dimension, compute a 1D array of length + # `num_traces` giving each trace's integer coordinate along that axis (via np.searchsorted). + # Cast to uint32. + idx_arrays: list[np.ndarray] = [] + for dim in self.dims[:-1]: + hdr_vals = index_headers[dim.name] # shape: (num_traces,) + coords = np.searchsorted(dim, hdr_vals) # integer indices + coords = coords.astype(np.uint32) + idx_arrays.append(coords) + + # Store as a tuple so that header_index_arrays[d][i] is "trace i's index along axis d" + self.header_index_arrays = tuple(idx_arrays) + + # We no longer allocate `self.map` or `self.live_mask` here. + # The full grid shape is `self.shape`, but mapping is done lazily per chunk. + return + + def get_traces_for_chunk(self, chunk_slices: tuple[slice, ...]) -> np.ndarray: + """Return all trace IDs whose grid-coordinates fall inside the given chunk slices. + + Args: + chunk_slices: Tuple of slice objects, one per grid dimension. For example, + (slice(i0, i1), slice(j0, j1), ...) corresponds to a single Zarr chunk + in index space (excluding the sample axis). + + Returns: + A 1D NumPy array of trace indices (0-based) that lie within the hyper-rectangle defined + by `chunk_slices`. If no traces fall in this chunk, returns an empty array. """ - # Determine data type for map based on grid size - grid_size = np.prod(self.shape[:-1], dtype=np.uint64) - map_dtype = np.uint64 if grid_size > UINT32_MAX else np.uint32 - fill_value = np.iinfo(map_dtype).max - - # Initialize Zarr arrays - live_shape = self.shape[:-1] - chunks = get_constrained_chunksize( - shape=live_shape, - dtype=map_dtype, - max_bytes=self._INTERNAL_CHUNK_SIZE_TARGET, - ) - self.map = zarr.full(live_shape, fill_value, dtype=map_dtype, chunks=chunks) - self.live_mask = zarr.zeros(live_shape, dtype=bool, chunks=chunks) - - # Calculate batch size - memory_per_trace_index = index_headers.itemsize - batch_size = max(1, int(self._TARGET_MEMORY_PER_BATCH / memory_per_trace_index)) - total_live_traces = index_headers.size - - # Process headers in batches - for start in range(0, total_live_traces, batch_size): - end = min(start + batch_size, total_live_traces) - live_dim_indices = [] - - # Compute indices for the batch - for dim in self.dims[:-1]: - dim_hdr = index_headers[dim.name][start:end] - indices = np.searchsorted(dim, dim_hdr).astype(np.uint32) - live_dim_indices.append(indices) - live_dim_indices = tuple(live_dim_indices) - - # Assign trace indices - trace_indices = np.arange(start, end, dtype=np.uint64) - self.map.vindex[live_dim_indices] = trace_indices - self.live_mask.vindex[live_dim_indices] = True + # Initialize a boolean mask over all traces (shape: (num_traces,)) + mask = np.ones((self.num_traces,), dtype=bool) + + for dim_idx, sl in enumerate(chunk_slices): + arr = self.header_index_arrays[dim_idx] # shape: (num_traces,) + start, stop = sl.start, sl.stop + if start is not None: + mask &= (arr >= start) + if stop is not None: + mask &= (arr < stop) + if not mask.any(): + # No traces remain after this dimension's filtering + return np.empty((0,), dtype=np.uint32) + + # Gather the trace IDs that survived all dimension tests + trace_ids = np.nonzero(mask)[0].astype(np.uint32) + return trace_ids class GridSerializer(Serializer): diff --git a/src/mdio/segy/_workers.py b/src/mdio/segy/_workers.py index 301847a6c..c0295dc22 100644 --- a/src/mdio/segy/_workers.py +++ b/src/mdio/segy/_workers.py @@ -77,45 +77,53 @@ def trace_worker( Returns: Partial statistics for chunk, or None """ - # Special case where there are no traces inside chunk. - live_subset = grid.live_mask[chunk_indices[:-1]] - - if np.count_nonzero(live_subset) == 0: + from time import time + start_time = time() + # Determine which trace IDs fall into this chunk + trace_ids = grid.get_traces_for_chunk(chunk_indices[:-1]) + if trace_ids.size == 0: return None - # Let's get trace numbers from grid map using the chunk indices. - seq_trace_indices = grid.map[chunk_indices[:-1]] - - tmp_data = np.zeros(seq_trace_indices.shape + (grid.shape[-1],), dtype=data_array.dtype) - tmp_metadata = np.zeros(seq_trace_indices.shape, dtype=metadata_array.dtype) - - del grid # To save some memory - - # Read headers and traces for block - valid_indices = seq_trace_indices[live_subset] - - traces = segy_file.trace[valid_indices.tolist()] + # Read headers and traces for the selected trace IDs + traces = segy_file.trace[trace_ids.tolist()] headers, samples = traces["header"], traces["data"] - tmp_metadata[live_subset] = headers.view(tmp_metadata.dtype) - tmp_data[live_subset] = samples - - # Flush metadata to zarr + # Build a temporary buffer for data and metadata for this chunk + chunk_shape = tuple(sli.stop - sli.start for sli in chunk_indices[:-1]) + (grid.shape[-1],) + tmp_data = np.zeros(chunk_shape, dtype=data_array.dtype) + meta_shape = tuple(sli.stop - sli.start for sli in chunk_indices[:-1]) + tmp_metadata = np.zeros(meta_shape, dtype=metadata_array.dtype) + + # Compute local coordinates within the chunk for each trace + local_coords: list[np.ndarray] = [] + for dim_idx, sl in enumerate(chunk_indices[:-1]): + hdr_arr = grid.header_index_arrays[dim_idx] + local_idx = (hdr_arr[trace_ids] - sl.start).astype(int) + local_coords.append(local_idx) + full_idx = tuple(local_coords) + (slice(None),) + + # Populate the temporary buffers + tmp_data[full_idx] = samples + tmp_metadata[tuple(local_coords)] = headers.view(tmp_metadata.dtype) + + # Flush metadata to Zarr metadata_array.set_basic_selection(selection=chunk_indices[:-1], value=tmp_metadata) + # Determine nonzero samples and early-exit if none nonzero_mask = samples != 0 - nonzero_count = nonzero_mask.sum(dtype="uint32") - + nonzero_count = int(nonzero_mask.sum()) if nonzero_count == 0: return None + # Flush data to Zarr data_array.set_basic_selection(selection=chunk_indices, value=tmp_data) # Calculate statistics - tmp_data = samples[nonzero_mask] - chunk_sum = tmp_data.sum(dtype="float64") - chunk_sum_squares = np.square(tmp_data, dtype="float64").sum() - min_val = tmp_data.min() - max_val = tmp_data.max() - - return nonzero_count, chunk_sum, chunk_sum_squares, min_val, max_val + flattened_nonzero = samples[nonzero_mask] + chunk_sum = float(flattened_nonzero.sum(dtype="float64")) + chunk_sum_squares = float(np.square(flattened_nonzero, dtype="float64").sum()) + min_val = float(flattened_nonzero.min()) + max_val = float(flattened_nonzero.max()) + + print(f"Time taken: {time() - start_time} seconds") + return (nonzero_count, chunk_sum, chunk_sum_squares, min_val, max_val) diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py index 9a4b59bb4..d279308ec 100644 --- a/src/mdio/segy/blocked_io.py +++ b/src/mdio/segy/blocked_io.py @@ -1,5 +1,32 @@ """Functions for doing blocked I/O from SEG-Y.""" +# from __future__ import annotations + +# import multiprocessing as mp +# import os +# from concurrent.futures import ProcessPoolExecutor +# from itertools import repeat +# from typing import TYPE_CHECKING, Any + +# import numpy as np +# from psutil import cpu_count +# from tqdm.auto import tqdm +# import zarr + +# # from mdio.core.indexing import ChunkIterator +# # from mdio.segy._workers import trace_worker + +# from mdio.core.indexing import ChunkIterator +# from mdio.segy._workers import trace_worker +# from mdio.segy.creation import SegyPartRecord +# from mdio.segy.creation import concat_files +# from mdio.segy.creation import serialize_to_segy_stack +# from mdio.segy.utilities import find_trailing_ones_index + +# if TYPE_CHECKING: +# from segy import SegyFile +# from mdio.core import Grid + from __future__ import annotations import multiprocessing as mp @@ -22,6 +49,8 @@ from mdio.segy.creation import serialize_to_segy_stack from mdio.segy.utilities import find_trailing_ones_index +import zarr + if TYPE_CHECKING: from numpy.typing import NDArray from segy import SegyFactory @@ -32,40 +61,75 @@ default_cpus = cpu_count(logical=True) -def to_zarr(segy_file: SegyFile, grid: Grid, data_array: Array, header_array: Array) -> dict: +def _worker_reopen( + zarr_root_path: str, + data_var_path: str, + header_var_path: str, + segy_file: SegyFile, + grid: Grid, + chunk_indices: tuple[slice, ...], +) -> tuple[Any, ...] | None: + """ + Worker function that reopens the Zarr store in this process, + obtains fresh array handles, and calls the real trace_worker. + """ + root = zarr.open_group(zarr_root_path, mode="r+") + data_arr = root[data_var_path] + header_arr = root[header_var_path] + result = trace_worker(segy_file, data_arr, header_arr, grid, chunk_indices) + root.store.close() + return result + + +def to_zarr( + segy_file: SegyFile, + grid: Grid, + zarr_root_path: str, + data_var_path: str, + header_var_path: str, +) -> dict[str, Any]: """Blocked I/O from SEG-Y to chunked `zarr.core.Array`. + Each worker reopens the Zarr store independently to avoid lock contention when writing. + Args: segy_file: SEG-Y file instance. - grid: mdio.Grid instance - data_array: Handle for zarr.core.Array we are writing trace data - header_array: Handle for zarr.core.Array we are writing trace headers + grid: mdio.Grid instance. + zarr_root_path: Filesystem path (or URI) to the root of the MDIO Zarr store. + data_var_path: Path within the Zarr group for the data array (e.g., "data/chunked_012"). + header_var_path: Path within the Zarr group for the header array + (e.g., "metadata/chunked_012_trace_headers"). Returns: Global statistics for the SEG-Y as a dictionary. """ + # Open Zarr store only in the main process to retrieve shape/metadata + root = zarr.open_group(zarr_root_path, mode="r") + data_array_meta = root[data_var_path] # only for shape info # Initialize chunk iterator (returns next chunk slice indices each iteration) - chunker = ChunkIterator(data_array, chunk_samples=False) + chunker = ChunkIterator(data_array_meta, chunk_samples=False) num_chunks = len(chunker) + root.store.close() # close immediately - # 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) + # Determine number of workers + num_cpus_env = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus)) + num_workers = min(num_chunks, num_cpus_env) context = mp.get_context("spawn") - executor = ProcessPoolExecutor(max_workers=num_workers, mp_context=context) # Chunksize here is for multiprocessing, not Zarr chunksize. pool_chunksize, extra = divmod(num_chunks, num_workers * 4) - pool_chunksize += 1 if extra else pool_chunksize + pool_chunksize += 1 if extra else 0 tqdm_kw = {"unit": "block", "dynamic_ncols": True} - with executor: + + # Launch multiprocessing pool + with ProcessPoolExecutor(max_workers=num_workers, mp_context=context) as executor: lazy_work = executor.map( - trace_worker, # fn + _worker_reopen, + repeat(zarr_root_path), + repeat(data_var_path), + repeat(header_var_path), repeat(segy_file), - repeat(data_array), - repeat(header_array), repeat(grid), chunker, chunksize=pool_chunksize, @@ -78,38 +142,31 @@ def to_zarr(segy_file: SegyFile, grid: Grid, data_array: Array, header_array: Ar **tqdm_kw, ) - # This executes the lazy work. chunk_stats = list(lazy_work) - # This comes in as n_chunk x 5 columns. - # Columns in order: count, sum, sum of squared, min, max. - # We can compute global mean, std, rms, min, max. - # Transposing because we want each statistic as a row to unpack later. - # REF: https://math.stackexchange.com/questions/1547141/aggregating-standard-deviation-to-a-summary-point # noqa: E501 - # REF: https://www.mathwords.com/r/root_mean_square.htm + # Aggregate statistics chunk_stats = [stat for stat in chunk_stats if stat is not None] + # Each stat: (count, sum, sum_sq, min, max). Transpose to unpack rows. + glob_count, glob_sum, glob_sum_square, glob_min, glob_max = zip(*chunk_stats) - chunk_stats = zip(*chunk_stats) # noqa: B905 - glob_count, glob_sum, glob_sum_square, glob_min, glob_max = chunk_stats - - glob_count = np.sum(glob_count) # Comes in as `uint32` - glob_sum = np.sum(glob_sum) # `float64` - glob_sum_square = np.sum(glob_sum_square) # `float64` - glob_min = np.min(glob_min) # `float32` - glob_max = np.max(glob_max) # `float32` + glob_count = np.sum(np.array(glob_count, dtype=np.uint64)) + glob_sum = np.sum(np.array(glob_sum, dtype=np.float64)) + glob_sum_square = np.sum(np.array(glob_sum_square, dtype=np.float64)) + glob_min = np.min(np.array(glob_min, dtype=np.float32)) + glob_max = np.max(np.array(glob_max, dtype=np.float32)) glob_mean = glob_sum / glob_count glob_std = np.sqrt(glob_sum_square / glob_count - (glob_sum / glob_count) ** 2) glob_rms = np.sqrt(glob_sum_square / glob_count) - # We need to write these as float64 because float32 is not JSON serializable - # Trace data is originally float32, hence min/max - glob_min = glob_min.min().astype("float64") - glob_max = glob_max.max().astype("float64") + # Convert to float64 for JSON compatibility + glob_min = float(glob_min) + glob_max = float(glob_max) return {"mean": glob_mean, "std": glob_std, "rms": glob_rms, "min": glob_min, "max": glob_max} + def segy_record_concat( block_records: NDArray, file_root: str, From 30a64724ccf3103b8d8fde241fd0f85ead6f96bc Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Mon, 2 Jun 2025 16:35:53 +0000 Subject: [PATCH 02/16] Remove prints --- src/mdio/converters/segy.py | 14 -------------- src/mdio/segy/_workers.py | 3 --- 2 files changed, 17 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 7e821fc76..9f2958c8a 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -369,7 +369,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 storage_options_input = storage_options_input or {} storage_options_output = storage_options_output or {} - print("pre-setup") # Open SEG-Y with MDIO's SegySpec. Endianness will be inferred. mdio_spec = mdio_segy_spec() segy_settings = SegySettings(storage_options=storage_options_input) @@ -379,7 +378,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 binary_header = segy.binary_header num_traces = segy.num_traces - print("pre-index") # Index the dataset using a spec that interprets the user provided index headers. index_fields: list[HeaderField] = [] for name, byte, format_ in zip(index_names, index_bytes, index_types, strict=True): @@ -387,7 +385,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 mdio_spec_grid = mdio_spec.customize(trace_header_fields=index_fields) segy_grid = SegyFile(url=segy_path, spec=mdio_spec_grid, settings=segy_settings) - print("pre-get_grid_plan") dimensions, chunksize, index_headers = get_grid_plan( segy_file=segy_grid, return_headers=True, @@ -395,12 +392,9 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 grid_overrides=grid_overrides, ) grid = Grid(dims=dimensions) - print("pre-grid_density_qc") grid_density_qc(grid, num_traces) - print("pre-build_map") grid.build_map(index_headers) - print("pre-valid_mask") # Check grid validity by ensuring every trace's header-index is within dimension bounds valid_mask = np.ones(grid.num_traces, dtype=bool) for d_idx in range(len(grid.header_index_arrays)): @@ -415,7 +409,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 logger.warning("Ingestion grid shape: %s.", grid.shape) raise GridTraceCountError(valid_count, num_traces) - print("pre-chunksize") if chunksize is None: dim_count = len(index_names) + 1 if dim_count == 2: # noqa: PLR2004 @@ -438,7 +431,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None] suffix = "".join(suffix) - print("pre-compressors") compressors = get_compressor(lossless, compression_tolerance) header_dtype = segy.spec.trace.header.dtype.newbyteorder("=") var_conf = MDIOVariableConfig( @@ -450,7 +442,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ) config = MDIOCreateConfig(path=mdio_path_or_buffer, grid=grid, variables=[var_conf]) - print("pre-create_empty") root_group = create_empty( config, overwrite=overwrite, @@ -462,7 +453,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 data_array = data_group[f"chunked_{suffix}"] header_array = meta_group[f"chunked_{suffix}_trace_headers"] - print("pre-live_mask") live_mask_array = meta_group["live_mask"] # 'live_mask_array' has the same first N–1 dims as 'grid.shape[:-1]' # Build a ChunkIterator over the live_mask (no sample axis) @@ -494,12 +484,10 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 nonzero_count = grid.num_traces - print("pre-write_attribute") write_attribute(name="trace_count", zarr_group=root_group, attribute=nonzero_count) write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n")) write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) - print("pre-to_zarr") # Write traces zarr_root = mdio_path_or_buffer # the same path you passed earlier to create_empty data_var = f"data/chunked_{suffix}" @@ -513,10 +501,8 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 header_var_path=header_var, ) - print("pre-write_attribute") # Write actual stats for key, value in stats.items(): write_attribute(name=key, zarr_group=root_group, attribute=value) - print("pre-consolidate_metadata") zarr.consolidate_metadata(root_group.store) \ No newline at end of file diff --git a/src/mdio/segy/_workers.py b/src/mdio/segy/_workers.py index c0295dc22..db8a54f11 100644 --- a/src/mdio/segy/_workers.py +++ b/src/mdio/segy/_workers.py @@ -77,8 +77,6 @@ def trace_worker( Returns: Partial statistics for chunk, or None """ - from time import time - start_time = time() # Determine which trace IDs fall into this chunk trace_ids = grid.get_traces_for_chunk(chunk_indices[:-1]) if trace_ids.size == 0: @@ -125,5 +123,4 @@ def trace_worker( min_val = float(flattened_nonzero.min()) max_val = float(flattened_nonzero.max()) - print(f"Time taken: {time() - start_time} seconds") return (nonzero_count, chunk_sum, chunk_sum_squares, min_val, max_val) From a23fb8410a47985d7c9cc12255dcef305a739165 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Mon, 2 Jun 2025 20:41:32 +0000 Subject: [PATCH 03/16] Fix to_zarr function signature to not be changed, don't reopen the zarr store --- src/mdio/converters/segy.py | 12 ++---- src/mdio/segy/blocked_io.py | 76 +++++-------------------------------- 2 files changed, 13 insertions(+), 75 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 9f2958c8a..aed169f0d 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -354,7 +354,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ... grid_overrides={"HasDuplicates": True}, ... ) """ - print("Entering segy_to_mdio") index_names = index_names or [f"dim_{i}" for i in range(len(index_bytes))] index_types = index_types or ["int32"] * len(index_bytes) @@ -379,7 +378,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 num_traces = segy.num_traces # Index the dataset using a spec that interprets the user provided index headers. - index_fields: list[HeaderField] = [] + index_fields = [] for name, byte, format_ in zip(index_names, index_bytes, index_types, strict=True): index_fields.append(HeaderField(name=name, byte=byte, format=format_)) mdio_spec_grid = mdio_spec.customize(trace_header_fields=index_fields) @@ -489,16 +488,11 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) # Write traces - zarr_root = mdio_path_or_buffer # the same path you passed earlier to create_empty - data_var = f"data/chunked_{suffix}" - header_var = f"metadata/chunked_{suffix}_trace_headers" - stats = blocked_io.to_zarr( segy_file=segy, grid=grid, - zarr_root_path=zarr_root, - data_var_path=data_var, - header_var_path=header_var, + data_array=data_array, + header_array=header_array, ) # Write actual stats diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py index d279308ec..7af1c5f18 100644 --- a/src/mdio/segy/blocked_io.py +++ b/src/mdio/segy/blocked_io.py @@ -1,32 +1,5 @@ """Functions for doing blocked I/O from SEG-Y.""" -# from __future__ import annotations - -# import multiprocessing as mp -# import os -# from concurrent.futures import ProcessPoolExecutor -# from itertools import repeat -# from typing import TYPE_CHECKING, Any - -# import numpy as np -# from psutil import cpu_count -# from tqdm.auto import tqdm -# import zarr - -# # from mdio.core.indexing import ChunkIterator -# # from mdio.segy._workers import trace_worker - -# from mdio.core.indexing import ChunkIterator -# from mdio.segy._workers import trace_worker -# from mdio.segy.creation import SegyPartRecord -# from mdio.segy.creation import concat_files -# from mdio.segy.creation import serialize_to_segy_stack -# from mdio.segy.utilities import find_trailing_ones_index - -# if TYPE_CHECKING: -# from segy import SegyFile -# from mdio.core import Grid - from __future__ import annotations import multiprocessing as mp @@ -35,6 +8,7 @@ from itertools import repeat from pathlib import Path from typing import TYPE_CHECKING +from typing import Any import numpy as np from dask.array import Array @@ -55,61 +29,33 @@ from numpy.typing import NDArray from segy import SegyFactory from segy import SegyFile + from zarr import Array as ZarrArray from mdio.core import Grid default_cpus = cpu_count(logical=True) -def _worker_reopen( - zarr_root_path: str, - data_var_path: str, - header_var_path: str, - segy_file: SegyFile, - grid: Grid, - chunk_indices: tuple[slice, ...], -) -> tuple[Any, ...] | None: - """ - Worker function that reopens the Zarr store in this process, - obtains fresh array handles, and calls the real trace_worker. - """ - root = zarr.open_group(zarr_root_path, mode="r+") - data_arr = root[data_var_path] - header_arr = root[header_var_path] - result = trace_worker(segy_file, data_arr, header_arr, grid, chunk_indices) - root.store.close() - return result - - def to_zarr( segy_file: SegyFile, grid: Grid, - zarr_root_path: str, - data_var_path: str, - header_var_path: str, + data_array: ZarrArray, + header_array: ZarrArray, ) -> dict[str, Any]: """Blocked I/O from SEG-Y to chunked `zarr.core.Array`. - Each worker reopens the Zarr store independently to avoid lock contention when writing. - Args: segy_file: SEG-Y file instance. grid: mdio.Grid instance. - zarr_root_path: Filesystem path (or URI) to the root of the MDIO Zarr store. - data_var_path: Path within the Zarr group for the data array (e.g., "data/chunked_012"). - header_var_path: Path within the Zarr group for the header array - (e.g., "metadata/chunked_012_trace_headers"). + data_array: Zarr array for storing trace data. + header_array: Zarr array for storing trace headers. Returns: Global statistics for the SEG-Y as a dictionary. """ - # Open Zarr store only in the main process to retrieve shape/metadata - root = zarr.open_group(zarr_root_path, mode="r") - data_array_meta = root[data_var_path] # only for shape info # Initialize chunk iterator (returns next chunk slice indices each iteration) - chunker = ChunkIterator(data_array_meta, chunk_samples=False) + chunker = ChunkIterator(data_array, chunk_samples=False) num_chunks = len(chunker) - root.store.close() # close immediately # Determine number of workers num_cpus_env = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus)) @@ -125,11 +71,10 @@ def to_zarr( # Launch multiprocessing pool with ProcessPoolExecutor(max_workers=num_workers, mp_context=context) as executor: lazy_work = executor.map( - _worker_reopen, - repeat(zarr_root_path), - repeat(data_var_path), - repeat(header_var_path), + trace_worker, repeat(segy_file), + repeat(data_array), + repeat(header_array), repeat(grid), chunker, chunksize=pool_chunksize, @@ -166,7 +111,6 @@ def to_zarr( return {"mean": glob_mean, "std": glob_std, "rms": glob_rms, "min": glob_min, "max": glob_max} - def segy_record_concat( block_records: NDArray, file_root: str, From 7bf7767da2f266331e5bc82015d4fb2598c86170 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 14:03:16 +0000 Subject: [PATCH 04/16] Optimization to try and add print for profiling --- src/mdio/converters/segy.py | 13 ++++++++++++- src/mdio/segy/_workers.py | 9 ++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index aed169f0d..ec106bd39 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -472,7 +472,14 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 local_coords: list[np.ndarray] = [] for dim_idx, sl in enumerate(chunk_indices): hdr_arr = grid.header_index_arrays[dim_idx] - local_idx = (hdr_arr[trace_ids] - sl.start).astype(int) + # Optimize memory usage: hdr_arr and trace_ids are already uint32, + # sl.start is int, so result should naturally be int32/uint32. + # Avoid unnecessary astype conversion to int64. + indexed_coords = hdr_arr[trace_ids] # uint32 array + local_idx = indexed_coords - sl.start # remains uint32 + # Only convert dtype if necessary for indexing (numpy requires int for indexing) + if local_idx.dtype != np.intp: + local_idx = local_idx.astype(np.intp) local_coords.append(local_idx) # Mark live cells in the temporary block @@ -487,6 +494,10 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n")) write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) + from datetime import datetime + + print("The livemask was written at time:", datetime.now().strftime("%H:%M:%S")) + # Write traces stats = blocked_io.to_zarr( segy_file=segy, diff --git a/src/mdio/segy/_workers.py b/src/mdio/segy/_workers.py index db8a54f11..efed69071 100644 --- a/src/mdio/segy/_workers.py +++ b/src/mdio/segy/_workers.py @@ -96,7 +96,14 @@ def trace_worker( local_coords: list[np.ndarray] = [] for dim_idx, sl in enumerate(chunk_indices[:-1]): hdr_arr = grid.header_index_arrays[dim_idx] - local_idx = (hdr_arr[trace_ids] - sl.start).astype(int) + # Optimize memory usage: hdr_arr and trace_ids are already uint32, + # sl.start is int, so result should naturally be int32/uint32. + # Avoid unnecessary astype conversion to int64. + indexed_coords = hdr_arr[trace_ids] # uint32 array + local_idx = indexed_coords - sl.start # remains uint32 + # Only convert dtype if necessary for indexing (numpy requires int for indexing) + if local_idx.dtype != np.intp: + local_idx = local_idx.astype(np.intp) local_coords.append(local_idx) full_idx = tuple(local_coords) + (slice(None),) From 942c27ee26b2158a0100be68a83947c7c3ba910b Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 14:20:46 +0000 Subject: [PATCH 05/16] Enhance profiling prints --- src/mdio/converters/segy.py | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index ec106bd39..ce003b40c 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -354,6 +354,12 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ... grid_overrides={"HasDuplicates": True}, ... ) """ + from datetime import datetime + import time + + start_time = datetime.fromtimestamp(time.time()) + print(f"Starting SEG-Y to MDIO conversion at: {start_time.strftime('%H:%M:%S.%f')}") + index_names = index_names or [f"dim_{i}" for i in range(len(index_bytes))] index_types = index_types or ["int32"] * len(index_bytes) @@ -377,6 +383,8 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 binary_header = segy.binary_header num_traces = segy.num_traces + print(f"SEG-Y file opened with {num_traces} traces at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") + # Index the dataset using a spec that interprets the user provided index headers. index_fields = [] for name, byte, format_ in zip(index_names, index_bytes, index_types, strict=True): @@ -394,6 +402,8 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 grid_density_qc(grid, num_traces) grid.build_map(index_headers) + print(f"Grid built and mapped at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") + # Check grid validity by ensuring every trace's header-index is within dimension bounds valid_mask = np.ones(grid.num_traces, dtype=bool) for d_idx in range(len(grid.header_index_arrays)): @@ -441,6 +451,8 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ) config = MDIOCreateConfig(path=mdio_path_or_buffer, grid=grid, variables=[var_conf]) + print(f"Starting MDIO file creation at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") + root_group = create_empty( config, overwrite=overwrite, @@ -457,6 +469,8 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 # Build a ChunkIterator over the live_mask (no sample axis) from mdio.core.indexing import ChunkIterator + print(f"Starting live mask creation at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") + chunker = ChunkIterator(live_mask_array, chunk_samples=False) for chunk_indices in chunker: # chunk_indices is a tuple of N–1 slice objects @@ -494,9 +508,10 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n")) write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) - from datetime import datetime + local_time = datetime.fromtimestamp(time.time()) + print(f"The livemask was written at time: {local_time.strftime('%H:%M:%S.%f')}") - print("The livemask was written at time:", datetime.now().strftime("%H:%M:%S")) + print(f"Starting trace writing at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") # Write traces stats = blocked_io.to_zarr( @@ -510,4 +525,9 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 for key, value in stats.items(): write_attribute(name=key, zarr_group=root_group, attribute=value) - zarr.consolidate_metadata(root_group.store) \ No newline at end of file + zarr.consolidate_metadata(root_group.store) + + end_time = datetime.fromtimestamp(time.time()) + print(f"SEG-Y to MDIO conversion completed at: {end_time.strftime('%H:%M:%S.%f')}") + duration = end_time - start_time + print(f"Total conversion time: {duration.total_seconds():.2f} seconds") \ No newline at end of file From 4274561b4eebab1894509d17d7520d74e5fa87d2 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 14:34:49 +0000 Subject: [PATCH 06/16] Investigate manual memory management --- src/mdio/converters/segy.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index ce003b40c..22c209d96 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -511,6 +511,16 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 local_time = datetime.fromtimestamp(time.time()) print(f"The livemask was written at time: {local_time.strftime('%H:%M:%S.%f')}") + # Clean up live mask related variables + del live_mask_array + del chunker + del block + del local_coords + del trace_ids + del chunk_indices + import gc + gc.collect() + print(f"Starting trace writing at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") # Write traces From 6cb0efe835a07d00ca8590a4abb2d72efb13d4fb Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 15:39:45 +0000 Subject: [PATCH 07/16] Cleanup profiling prints --- src/mdio/converters/segy.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 22c209d96..93230686c 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -354,12 +354,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ... grid_overrides={"HasDuplicates": True}, ... ) """ - from datetime import datetime - import time - - start_time = datetime.fromtimestamp(time.time()) - print(f"Starting SEG-Y to MDIO conversion at: {start_time.strftime('%H:%M:%S.%f')}") - index_names = index_names or [f"dim_{i}" for i in range(len(index_bytes))] index_types = index_types or ["int32"] * len(index_bytes) @@ -383,8 +377,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 binary_header = segy.binary_header num_traces = segy.num_traces - print(f"SEG-Y file opened with {num_traces} traces at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") - # Index the dataset using a spec that interprets the user provided index headers. index_fields = [] for name, byte, format_ in zip(index_names, index_bytes, index_types, strict=True): @@ -402,8 +394,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 grid_density_qc(grid, num_traces) grid.build_map(index_headers) - print(f"Grid built and mapped at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") - # Check grid validity by ensuring every trace's header-index is within dimension bounds valid_mask = np.ones(grid.num_traces, dtype=bool) for d_idx in range(len(grid.header_index_arrays)): @@ -451,8 +441,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 ) config = MDIOCreateConfig(path=mdio_path_or_buffer, grid=grid, variables=[var_conf]) - print(f"Starting MDIO file creation at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") - root_group = create_empty( config, overwrite=overwrite, @@ -469,8 +457,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 # Build a ChunkIterator over the live_mask (no sample axis) from mdio.core.indexing import ChunkIterator - print(f"Starting live mask creation at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") - chunker = ChunkIterator(live_mask_array, chunk_samples=False) for chunk_indices in chunker: # chunk_indices is a tuple of N–1 slice objects @@ -508,9 +494,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n")) write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) - local_time = datetime.fromtimestamp(time.time()) - print(f"The livemask was written at time: {local_time.strftime('%H:%M:%S.%f')}") - # Clean up live mask related variables del live_mask_array del chunker @@ -521,8 +504,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 import gc gc.collect() - print(f"Starting trace writing at: {datetime.fromtimestamp(time.time()).strftime('%H:%M:%S.%f')}") - # Write traces stats = blocked_io.to_zarr( segy_file=segy, @@ -536,8 +517,3 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 write_attribute(name=key, zarr_group=root_group, attribute=value) zarr.consolidate_metadata(root_group.store) - - end_time = datetime.fromtimestamp(time.time()) - print(f"SEG-Y to MDIO conversion completed at: {end_time.strftime('%H:%M:%S.%f')}") - duration = end_time - start_time - print(f"Total conversion time: {duration.total_seconds():.2f} seconds") \ No newline at end of file From 4500bfa688f00addd8d44e37ab5a974059609733 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 19:47:36 +0000 Subject: [PATCH 08/16] Fix live mask chunk iteration --- src/mdio/converters/segy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 93230686c..a431bd186 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -457,7 +457,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 # Build a ChunkIterator over the live_mask (no sample axis) from mdio.core.indexing import ChunkIterator - chunker = ChunkIterator(live_mask_array, chunk_samples=False) + chunker = ChunkIterator(live_mask_array, chunk_samples=True) for chunk_indices in chunker: # chunk_indices is a tuple of N–1 slice objects trace_ids = grid.get_traces_for_chunk(chunk_indices) From fa60f1425484b2055e68afd7c75f5a8c31655b7b Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 19:53:55 +0000 Subject: [PATCH 09/16] Revert accidentally deleted comments --- src/mdio/segy/blocked_io.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py index 7af1c5f18..d7a1dd4b6 100644 --- a/src/mdio/segy/blocked_io.py +++ b/src/mdio/segy/blocked_io.py @@ -89,6 +89,9 @@ def to_zarr( chunk_stats = list(lazy_work) + # Transposing because we want each statistic as a row to unpack later. + # REF: https://math.stackexchange.com/questions/1547141/aggregating-standard-deviation-to-a-summary-point # noqa: E501 + # REF: https://www.mathwords.com/r/root_mean_square.htm # Aggregate statistics chunk_stats = [stat for stat in chunk_stats if stat is not None] # Each stat: (count, sum, sum_sq, min, max). Transpose to unpack rows. From 4e35a13790798f1f6d6dc02c0c6f3a8bac45d4b8 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 19:57:25 +0000 Subject: [PATCH 10/16] Relocate context definition and add back comments --- src/mdio/segy/blocked_io.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py index d7a1dd4b6..57ce72e11 100644 --- a/src/mdio/segy/blocked_io.py +++ b/src/mdio/segy/blocked_io.py @@ -60,7 +60,6 @@ def to_zarr( # Determine number of workers num_cpus_env = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus)) num_workers = min(num_chunks, num_cpus_env) - context = mp.get_context("spawn") # Chunksize here is for multiprocessing, not Zarr chunksize. pool_chunksize, extra = divmod(num_chunks, num_workers * 4) @@ -68,6 +67,10 @@ def to_zarr( tqdm_kw = {"unit": "block", "dynamic_ncols": True} + # 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 + context = mp.get_context("spawn") + # Launch multiprocessing pool with ProcessPoolExecutor(max_workers=num_workers, mp_context=context) as executor: lazy_work = executor.map( From a7ab60ef090aeed43ea9ad720cb6a49667af65b7 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 3 Jun 2025 20:42:11 +0000 Subject: [PATCH 11/16] Memory management --- src/mdio/converters/segy.py | 37 +++++++++++++++++++++++++------------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index a431bd186..0c13e30f2 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -456,17 +456,19 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 # 'live_mask_array' has the same first N–1 dims as 'grid.shape[:-1]' # Build a ChunkIterator over the live_mask (no sample axis) from mdio.core.indexing import ChunkIterator + import gc chunker = ChunkIterator(live_mask_array, chunk_samples=True) for chunk_indices in chunker: # chunk_indices is a tuple of N–1 slice objects trace_ids = grid.get_traces_for_chunk(chunk_indices) if trace_ids.size == 0: + # Free memory immediately for empty chunks + del trace_ids continue # Build a temporary boolean block of shape = chunk shape - block_shape = tuple(sl.stop - sl.start for sl in chunk_indices) - block = np.zeros(block_shape, dtype=bool) + block = np.zeros(tuple(sl.stop - sl.start for sl in chunk_indices), dtype=bool) # Compute local coords within this block for each trace_id local_coords: list[np.ndarray] = [] @@ -477,16 +479,37 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 # Avoid unnecessary astype conversion to int64. indexed_coords = hdr_arr[trace_ids] # uint32 array local_idx = indexed_coords - sl.start # remains uint32 + # Free indexed_coords immediately + del indexed_coords + # Only convert dtype if necessary for indexing (numpy requires int for indexing) if local_idx.dtype != np.intp: local_idx = local_idx.astype(np.intp) local_coords.append(local_idx) + # local_idx is now owned by local_coords list, safe to continue + + # Free trace_ids as soon as we're done with it + del trace_ids # Mark live cells in the temporary block block[tuple(local_coords)] = True + + # Free local_coords immediately after use + del local_coords # Write the entire block to Zarr at once live_mask_array.set_basic_selection(selection=chunk_indices, value=block) + + # Free block immediately after writing + del block + + # Force garbage collection periodically to free memory aggressively + gc.collect() + + # Final cleanup + del live_mask_array + del chunker + gc.collect() nonzero_count = grid.num_traces @@ -494,16 +517,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n")) write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) - # Clean up live mask related variables - del live_mask_array - del chunker - del block - del local_coords - del trace_ids - del chunk_indices - import gc - gc.collect() - # Write traces stats = blocked_io.to_zarr( segy_file=segy, From 5d38dbbd9b1f45dc5b9001cb0fc4615f6753892b Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 4 Jun 2025 14:13:18 +0000 Subject: [PATCH 12/16] Eager free of valid_mask --- src/mdio/converters/segy.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 0c13e30f2..587da6ef7 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -408,6 +408,10 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 logger.warning("Ingestion grid shape: %s.", grid.shape) raise GridTraceCountError(valid_count, num_traces) + import gc + del valid_mask + gc.collect() + if chunksize is None: dim_count = len(index_names) + 1 if dim_count == 2: # noqa: PLR2004 @@ -456,7 +460,6 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 # 'live_mask_array' has the same first N–1 dims as 'grid.shape[:-1]' # Build a ChunkIterator over the live_mask (no sample axis) from mdio.core.indexing import ChunkIterator - import gc chunker = ChunkIterator(live_mask_array, chunk_samples=True) for chunk_indices in chunker: From f58593bdce1c557f14b1b541c76ce273ae04590f Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 4 Jun 2025 14:20:03 +0000 Subject: [PATCH 13/16] Linting --- src/mdio/converters/segy.py | 13 +++++++------ src/mdio/core/grid.py | 20 ++++++++------------ src/mdio/segy/blocked_io.py | 4 +--- 3 files changed, 16 insertions(+), 21 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 587da6ef7..409bf0dcb 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -131,7 +131,7 @@ def get_compressor(lossless: bool, compression_tolerance: float = -1) -> Blosc | return compressor -def segy_to_mdio( # noqa: PLR0913, PLR0915 +def segy_to_mdio( # noqa: PLR0913, PLR0915, PLR0912 segy_path: str | Path, mdio_path_or_buffer: str | Path, index_bytes: Sequence[int], @@ -398,7 +398,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 valid_mask = np.ones(grid.num_traces, dtype=bool) for d_idx in range(len(grid.header_index_arrays)): coords = grid.header_index_arrays[d_idx] - valid_mask &= (coords < grid.shape[d_idx]) + valid_mask &= coords < grid.shape[d_idx] valid_count = int(np.count_nonzero(valid_mask)) if valid_count != num_traces: for dim_name in grid.dim_names: @@ -409,6 +409,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 raise GridTraceCountError(valid_count, num_traces) import gc + del valid_mask gc.collect() @@ -484,7 +485,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 local_idx = indexed_coords - sl.start # remains uint32 # Free indexed_coords immediately del indexed_coords - + # Only convert dtype if necessary for indexing (numpy requires int for indexing) if local_idx.dtype != np.intp: local_idx = local_idx.astype(np.intp) @@ -496,16 +497,16 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915 # Mark live cells in the temporary block block[tuple(local_coords)] = True - + # Free local_coords immediately after use del local_coords # Write the entire block to Zarr at once live_mask_array.set_basic_selection(selection=chunk_indices, value=block) - + # Free block immediately after writing del block - + # Force garbage collection periodically to free memory aggressively gc.collect() diff --git a/src/mdio/core/grid.py b/src/mdio/core/grid.py index 9ed9378c1..8bdc42ad5 100644 --- a/src/mdio/core/grid.py +++ b/src/mdio/core/grid.py @@ -7,14 +7,12 @@ from typing import TYPE_CHECKING import numpy as np -import zarr -from mdio.constants import UINT32_MAX from mdio.core import Dimension from mdio.core.serialization import Serializer -from mdio.core.utils_write import get_constrained_chunksize if TYPE_CHECKING: + import zarr from segy.arrays import HeaderArray from zarr import Array as ZarrArray @@ -112,8 +110,8 @@ def build_map(self, index_headers: HeaderArray) -> None: """Compute per-trace grid coordinates (lazy map). Instead of allocating a full `self.map` and `self.live_mask`, this computes, for each trace, - its integer index along each dimension (excluding the final sample dimension) and stores them in - `self.header_index_arrays`. The full mapping can then be derived chunk-by-chunk when writing. + its integer index along each dimension (excluding the sample dimension) and stores them in + `self.header_index_arrays`. The full mapping can then be derived chunkwise when writing. Args: index_headers: Header array containing dimension indices (length = number of traces). @@ -126,8 +124,8 @@ def build_map(self, index_headers: HeaderArray) -> None: # Cast to uint32. idx_arrays: list[np.ndarray] = [] for dim in self.dims[:-1]: - hdr_vals = index_headers[dim.name] # shape: (num_traces,) - coords = np.searchsorted(dim, hdr_vals) # integer indices + hdr_vals = index_headers[dim.name] # shape: (num_traces,) + coords = np.searchsorted(dim, hdr_vals) # integer indices coords = coords.astype(np.uint32) idx_arrays.append(coords) @@ -136,7 +134,6 @@ def build_map(self, index_headers: HeaderArray) -> None: # We no longer allocate `self.map` or `self.live_mask` here. # The full grid shape is `self.shape`, but mapping is done lazily per chunk. - return def get_traces_for_chunk(self, chunk_slices: tuple[slice, ...]) -> np.ndarray: """Return all trace IDs whose grid-coordinates fall inside the given chunk slices. @@ -157,16 +154,15 @@ def get_traces_for_chunk(self, chunk_slices: tuple[slice, ...]) -> np.ndarray: arr = self.header_index_arrays[dim_idx] # shape: (num_traces,) start, stop = sl.start, sl.stop if start is not None: - mask &= (arr >= start) + mask &= arr >= start if stop is not None: - mask &= (arr < stop) + mask &= arr < stop if not mask.any(): # No traces remain after this dimension's filtering return np.empty((0,), dtype=np.uint32) # Gather the trace IDs that survived all dimension tests - trace_ids = np.nonzero(mask)[0].astype(np.uint32) - return trace_ids + return np.nonzero(mask)[0].astype(np.uint32) class GridSerializer(Serializer): diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py index 57ce72e11..0498c5e5c 100644 --- a/src/mdio/segy/blocked_io.py +++ b/src/mdio/segy/blocked_io.py @@ -23,8 +23,6 @@ from mdio.segy.creation import serialize_to_segy_stack from mdio.segy.utilities import find_trailing_ones_index -import zarr - if TYPE_CHECKING: from numpy.typing import NDArray from segy import SegyFactory @@ -98,7 +96,7 @@ def to_zarr( # Aggregate statistics chunk_stats = [stat for stat in chunk_stats if stat is not None] # Each stat: (count, sum, sum_sq, min, max). Transpose to unpack rows. - glob_count, glob_sum, glob_sum_square, glob_min, glob_max = zip(*chunk_stats) + glob_count, glob_sum, glob_sum_square, glob_min, glob_max = zip(*chunk_stats, strict=False) glob_count = np.sum(np.array(glob_count, dtype=np.uint64)) glob_sum = np.sum(np.array(glob_sum, dtype=np.float64)) From 61fd99d2c4eb0adae16e69acf99a965396f9707f Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 4 Jun 2025 15:35:42 +0000 Subject: [PATCH 14/16] Fix pool chunking logic --- src/mdio/segy/blocked_io.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py index 0498c5e5c..3211c2029 100644 --- a/src/mdio/segy/blocked_io.py +++ b/src/mdio/segy/blocked_io.py @@ -61,7 +61,8 @@ def to_zarr( # Chunksize here is for multiprocessing, not Zarr chunksize. pool_chunksize, extra = divmod(num_chunks, num_workers * 4) - pool_chunksize += 1 if extra else 0 + if extra: + pool_chunksize += 1 tqdm_kw = {"unit": "block", "dynamic_ncols": True} From d4deac7f1ace174a36e749a6c66d26e849bf8f3f Mon Sep 17 00:00:00 2001 From: Brian Michell Date: Fri, 6 Jun 2025 13:50:29 -0500 Subject: [PATCH 15/16] Update to honor cpu count env var --- src/mdio/converters/segy.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index 409bf0dcb..fba958abe 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -354,6 +354,11 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915, PLR0912 ... grid_overrides={"HasDuplicates": True}, ... ) """ + from zarr.core.config import config as zarr_config + import os + num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", 1)) + zarr_config.set({"threading.max_workers": num_cpus}) + index_names = index_names or [f"dim_{i}" for i in range(len(index_bytes))] index_types = index_types or ["int32"] * len(index_bytes) @@ -521,6 +526,7 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915, PLR0912 write_attribute(name="text_header", zarr_group=meta_group, attribute=text_header.split("\n")) write_attribute(name="binary_header", zarr_group=meta_group, attribute=binary_header.to_dict()) + zarr_config.set({"threading.max_workers": 1}) # Write traces stats = blocked_io.to_zarr( segy_file=segy, From 1532d671ff9a052dd55e10f75401f9113fb46732 Mon Sep 17 00:00:00 2001 From: Brian Michell Date: Fri, 6 Jun 2025 14:27:05 -0500 Subject: [PATCH 16/16] Pre-commit --- src/mdio/converters/segy.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py index fba958abe..caae5793d 100644 --- a/src/mdio/converters/segy.py +++ b/src/mdio/converters/segy.py @@ -354,11 +354,13 @@ def segy_to_mdio( # noqa: PLR0913, PLR0915, PLR0912 ... grid_overrides={"HasDuplicates": True}, ... ) """ - from zarr.core.config import config as zarr_config import os - num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", 1)) + + from zarr.core.config import config as zarr_config + + num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", "1")) zarr_config.set({"threading.max_workers": num_cpus}) - + index_names = index_names or [f"dim_{i}" for i in range(len(index_bytes))] index_types = index_types or ["int32"] * len(index_bytes)