Optimize save_weights: rank-ordered layout with collective I/O (3)#84
Open
cattabiani wants to merge 13 commits into
Open
Optimize save_weights: rank-ordered layout with collective I/O (3)#84cattabiani wants to merge 13 commits into
cattabiani wants to merge 13 commits into
Conversation
Replace the per-node HDF5 write loop in _add_data with a single collective write per MPI rank. The file layout is now rank-ordered (each rank's segments contiguous) instead of globally GID-sorted. Changes: - _init_weights: keep all_cols in rank-concatenation order, build node_ids/offsets in rank order using groupby(sort=False) - _add_data: new 'start' parameter for optimized single-write path; legacy per-node path preserved for backward compat with unit tests - save_weights: use MPI_Scan for row offsets, all ranks participate in collective I/O (including empty ranks) - _write_neurite_types: same contiguous write pattern - MPI integration tests: order-independent comparison (per-node lookup) The SONATA spec does not require globally sorted node_ids. Readers (neurodamus) use np.where(node_ids == gid) for lookup.
Simplify _add_data to only the optimized single-block write path. The 'ids' parameter is removed; 'start' is now required. Unit tests updated to use start=0 (single-rank = offset 0). The 'backwards' test now verifies that data is written in coeffs column order (which matches the file layout in the new architecture).
GIDs are disjoint across ranks (round-robin), so no deduplication needed. Replace dict-based seen tracking with simple concatenation of each rank's np.unique output.
- Replace collective I/O with independent writes (non-overlapping slices) - Use fillvalue=1.0 instead of np.ones to avoid 17GB memory spike in init - Add log_rank0() utility for rank-0 progress messages - Add timing summary to save_weights - Warn if h5py lacks MPI support in multi-rank runs
cattabiani
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Replace the per-node HDF5 write loop in
_add_datawith a single collective write per MPI rank. Each rank now writes a contiguous block ofscaling_factorsin one operation.Fix: #83
Fix: #79
Fix: #73
Changes
_init_weights: File layout is now rank-ordered (rank 0's segments first, then rank 1's, etc.) instead of globally GID-sorted. Usesgroupby(sort=False)to preserve rank-concatenation order._add_data: Writescoeffs.values.Tas a single contiguous slice using h5py collective I/O. Falls back to a plain write when h5py lacks parallel support (single-rank / unit tests).save_weights: UsesMPI_Scan(comm.scan) to compute each rank's row offset. All ranks (including empty ones) participate in collective I/O. Emits a warning if running multi-rank without parallel h5py._write_neurite_types: Same single-write pattern with collective I/O fallback.Rationale
node_idsin electrode filesnp.where(node_ids == gid)— no sorted assumptionTesting
mpirun -n 2