fix(data): restore pre-allocated batch_ptr capacity on put

[lil-lon]

ALCHEMI Toolkit Pull Request

Description

Fixes #136: the first put() trims the pre-allocated batch_ptr in SegmentedLevelStorage, so every later put() fails the capacity checks and silently copies nothing, and zero() bakes in the trimmed size.

Root cause

batch_ptr must be tight for readers (len == num_graphs + 1; kernels derive counts from its shape) but must keep headroom for put() to append. The headroom's only record was batch_ptr.shape[0] itself, so the first trim destroyed it irreversibly — it cannot be re-derived from the data shape (elements, not segments).

Why it went unnoticed

GPUBuffer, the consumer of reusable Batch.empty buffers, re-extends batch_ptr after every write() from its externally-known capacity (_restore_batch_ptr_capacity, added in #4), which masked the bug on that path.

Fix

SegmentedLevelStorage records batch_ptr_capacity at construction and re-extends the trimmed batch_ptr just before the capacity checks in put() / compute_put_per_system_fit_mask().
Every Batch.empty buffer is covered at once, and the GPUBuffer restore becomes redundant and is removed.
The lazy-init/extend block in _ensure_buffer was already dead code: Batch.empty allocates batch_ptr eagerly with num_systems + 2 entries. Nothing else changes: the trim and the published tight batch_ptr are as before, and non-buffer storages record no capacity, so the restore is a no-op for them.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

#136

Changes Made

• SegmentedLevelStorage.__init__ records batch_ptr_capacity; new _restore_batch_ptr_capacity() re-extends a trimmed batch_ptr, called from put() and compute_put_per_system_fit_mask(); clone() carries the capacity
• Removed GPUBuffer._restore_batch_ptr_capacity and its call sites
• Regression tests: repeated put() into a Batch.empty buffer, put → zero → put reuse (the 🐛[BUG]: A pre-allocated Batch.empty buffer can only be filled once: every put() after the first is a silent no-op #136 repro), and a storage-level repeated-put test

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

I reported #136 and already had this fix in hand, so I am submitting it as a concrete reference. I am aware that direct code contributions are not currently being accepted during the public beta, and this touches a fairly deep data-layer class. So, please feel free to close this PR if that process still applies. I hope the diff and the regression tests are useful either way.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR fixes a silent data-loss bug (#136) in SegmentedLevelStorage where the first put() trimmed _batch_ptr to the number of active segments, destroying the headroom that Batch.empty had pre-allocated; every subsequent put() then silently copied nothing.

• SegmentedLevelStorage.__init__ now records _batch_ptr_capacity when the batch_ptr_capacity parameter is supplied, and the new _restore_batch_ptr_capacity() method re-extends the trimmed tensor (filling the tail with _batch_ptr[n-1], matching the constructor invariant) before the capacity checks in put() and compute_put_per_system_fit_mask().
• GPUBuffer._restore_batch_ptr_capacity and its call sites are removed as redundant; clone() is updated to carry _batch_ptr_capacity to clones.
• Three focused regression tests cover repeated put(), put → zero → put reuse, and storage-level append integrity.

Important Files Changed

Filename	Overview
nvalchemi/data/level_storage.py	Core fix: records _batch_ptr_capacity at construction and adds _restore_batch_ptr_capacity() called before capacity checks in put() and compute_put_per_system_fit_mask(). Logic is correct — tail-fill uses _batch_ptr[n-1] matching the __init__ invariant; clone() correctly propagates the capacity.
nvalchemi/dynamics/sinks.py	Removes the now-redundant GPUBuffer._restore_batch_ptr_capacity and its call sites; the PR description correctly identifies the _ensure_buffer block as dead code since Batch.empty already eagerly allocates _batch_ptr with the full capacity.
test/data/test_batch.py	Adds two well-targeted regression tests: repeated put() into an empty buffer, and the exact put → zero → put repro from #136. Both exercise the new _restore_batch_ptr_capacity path end-to-end.
test/data/test_level_storage.py	Adds a storage-level repeated-put test that validates data integrity (element values) after three successive puts, confirming correct append semantics at the SegmentedLevelStorage layer.
CHANGELOG.md	Changelog entry added under the correct section; accurately describes the root cause and fix.

_{Reviews (1): Last reviewed commit: "fix(data): restore pre-allocated batch_p..." | Re-trigger Greptile}

[lil-lon]

#141 fixes the same root cause (#136) by removing the trim in put(), so the level_storage.py change here becomes redundant once it merges.
Two parts of this PR stay relevant either way:

• the removal of the dead workaround in sinks.py
• the regression tests, including repeated put() without zero()

For the record, why this PR kept the trim: without it, every put leaves the target's batch_ptr inconsistent. Its length no longer matches the active segment count, and entries past the last written boundary keep stale values.
No current consumer reads that state, which is why #141 fixes #136.

Restoring capacity instead keeps "trimmed at rest" as a structural guarantee rather than a usage convention. That was the intent.

Happy to close this once #141 lands, or rebase it down to the sinks cleanup + tests if those are still useful.

[laserkelvin]

Thank you for the PR @lil-lon - since #141 came up in our own bug catching/fixing efforts, I'd prefer we merge that first, and then ask you to update your PR because I think we are still interested in the cleanup + tests you mentioned.

[lil-lon]

Thanks for taking a look!
Since #141 landed on 0.2.0-rc branch, I'll hold off until it is merged to main and then trim this PR down to the cleanup and tests.