497 ingestion memory by BrianMichell · Pull Request #558 · TGSAI/mdio-python

BrianMichell · 2025-06-03T15:42:41Z

Closes #497

Resolves memory usage issues by lazily computing indices instead of eagerly computing the full vindex array.

Before:

After:

…rr store

tasansal

i think we over complicated it and I don't understand what component contributes to reducing memory footprint. Some things that were compressed in memory are now full allocation etc. Can you please explain the algorithmic changes and maybe we can revert some that doesn't help with memory footprint

tasansal · 2025-06-03T15:51:50Z

+    # 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()
+


we don't necessarily have to do all this; can you profile it end-to-end without this and see if its necessary

Deeper dive into why I did this

Here's a zoomed in view of commit 4274561b4eebab1894509d17d7520d74e5fa87d2 (No eager del or gc)

Here's a zoomed in view after adding the eager garbage collection

These two plots show the period between from mdio.core.indexing import ChunkIterator and # Clean up live mask related variables in the current code.

It helped decrease heap allocation by ~200MiB to eagerly delete and garbage collect the objects.

*I failed to drop the cache on the After run of the initial PR message. These figures are both with dropped cache.

tasansal · 2025-06-03T15:54:07Z

+        block[tuple(local_coords)] = True
+
+        # Write the entire block to Zarr at once
+        live_mask_array.set_basic_selection(selection=chunk_indices, value=block)


what is the performance implication of this when writing to cloud storage? This will be A LOT of write requests in sequence. We could probably optimize the size of the block iteration (e.g. give a mock array to ChunkIterator with larger virtual chunks)

This is intended to operate on _TARGET_MEMORY_PER_BATCH sized live masks, which should end up with relatively few write requests.

That does bring up the issue that I seem to have accidentally nuked the internal logic to chunk it in the first place though.

Correction: I confused myself between the internal Grid chunking and the utils_write.py chunking size. This does leverage the persistent live_mask chunk size appropriately. block should be 512MiB as defined by MAX_SIZE_LIVE_MASK.

cool, makes sense then. can you validate how many write requests it does for a 100GB-sh file?

I've validated that it only reaches this block once per live mask chunk. For a 100GB file, this is once. Reduced the maximum size of the live mask and saw the appropriate amount of iterations.

tasansal · 2025-06-03T15:57:08Z

+    # 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)


Doesn't this use more memory? grid.live_mask is a compressed bool zarr array in memory, but we just allocated the full size with np.ones. Doesn't this increase memory consumption?

It does use more memory in the main process compared to what is live in 0.9.2 (below), but it has a lower peak memory consumption. This should only account for a few megabytes though. I'll look into any other low-hanging fruit as far as memory consumption.

it may become problematic if like mask is 4-5GB. We can probably still store it as a zarr array.

Ideally, valid_mask is freed immediately following the validation check (and I've eagerly collected it as well now). I don't believe it is necessary to hold the compressed mask at all.

tasansal · 2025-06-03T16:03:03Z

    # 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


this is interesting; what was the behaviour before?

Was it adding pool_chunksize if not extra?

Maybe simpler to

if not extra: pool_chunksize += 1

wdyt

The previous logic was adding pool_chunksize to itself in the negative branch of the ternary operator, which was causing tests to fail for some reason. The updated ternary made sure to not modify the variable in the negative case. I've refactored it to use an explicit if block so the logic is easier to understand at a glance.

tasansal · 2025-06-03T16:06:29Z

@@ -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:


i don't understand the need to change in this function and the next. the map and live_mask are compressed zarr arrays in memory and headers are already processed in batch.

This was actually the majority of the reason behind the change.

The vindex used to eagerly generate the map appeared to be responsible for the sharp spike around 11:05 in the initial figure. The intent of these changes was to effectively defer those computations to the distributed tasks.

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 # <--- This vindexing assignment self.live_mask.vindex[live_dim_indices] = True

hm, if i understood correctly, if we are writing the live mask from distributed tasks (which are chunked same as seismic) it will cause race conditions and corrupt the live mask (which has larger chunks). e.g. multiple workers may write to the same chunk at the same time. Have you checked if live mask is identical before/after this?

We are not writing the live mask in a distributed manner, and the Grid.live_mask chunking also differs from the live mask which is written to durable media. I believe this is 100% safe behavior.

tasansal · 2025-06-07T07:12:12Z

@BrianMichell can you please break this PR into components and benchmark the effect of memory for each significant change. I would like to avoid large code changes that may have minimal benefit.

BrianMichell · 2025-06-09T13:33:37Z

@BrianMichell can you please break this PR into components and benchmark the effect of memory for each significant change. I would like to avoid large code changes that may have minimal benefit.

Sure thing @tasansal. It's pretty difficult to break this into more atomic units for benchmarking, but I have these plots readily available.

The System memory plots may not be fully accurate of solely the MDIO ingestion process and it's spawned children. I made sure all system caches were dropped immediately before each run. The green plot should be f58593b

Not included in the above plot is applying only changes from d4deac7. This change still exhibited very high memory spikes during the grid building and first phase of MDIO writes before following the expected memory trend.

All time data here is singe-run, with only the green plot showing any performance regression.

BrianMichell added 7 commits June 2, 2025 16:15

Begin testing for lazy compute of grid

d27879f

Remove prints

30a6472

Fix to_zarr function signature to not be changed, don't reopen the za…

a23fb84

…rr store

Optimization to try and add print for profiling

7bf7767

Enhance profiling prints

942c27e

Investigate manual memory management

4274561

Cleanup profiling prints

6cb0efe

tasansal suggested changes Jun 3, 2025

View reviewed changes

BrianMichell added 9 commits June 3, 2025 19:47

Fix live mask chunk iteration

4500bfa

Revert accidentally deleted comments

fa60f14

Relocate context definition and add back comments

4e35a13

Memory management

a7ab60e

Eager free of valid_mask

5d38dbb

Linting

f58593b

Fix pool chunking logic

61fd99d

Update to honor cpu count env var

d4deac7

Pre-commit

1532d67

tasansal closed this Jun 26, 2025

		@@ -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:

Uh oh!

Conversation

BrianMichell commented Jun 3, 2025

Uh oh!

tasansal left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BrianMichell Jun 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tasansal commented Jun 7, 2025

Uh oh!

BrianMichell commented Jun 9, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

BrianMichell Jun 3, 2025 •

edited

Loading