perf(arrow): stream serialize-to-backend writes via BufferWritableBackend (LAB-766) - #247
Conversation
…kend (LAB-766) The write path materialized the whole serialized envelope before a single byte reached the backend — pyarrow's BufferOutputStream doubling-growth plus memory-pool retention held ~3.75x the payload, putting write peak RSS at ~5.6x logical for large DataFrames. Write-side twin of the get_buffer mmap read path (#171/#187): - BufferWritableBackend capability protocol (backends/base.py): writer-callback set_streaming(key, write_payload, ttl) against a seekable sink. A chunk iterator cannot express the Arrow envelope's checksum-first layout (the 8-byte xxHash3-64 precedes the IPC bytes it covers), so the producer hashes incrementally while streaming and patches the checksum in place — stored bytes stay byte-identical to the buffered path. - FileBackend.set_streaming: incremental write behind the existing temp-file + exclusive-flock + fsync + atomic-rename pattern; producer failures discard the partial write and leave any previous value intact; max_value_mb enforced on the on-disk size before rename. - ArrowSerializer.serialize_to_sink: streams the ~8 MiB record-batch IPC frames through a hashing tee straight into the sink, bypassing BufferOutputStream entirely; optional byte budget aborts over-limit values mid-stream (max_value_size, issue #163 ceiling preserved). - Write-flow gating: plaintext Arrow only (AES-256-GCM needs the whole ciphertext for its tag — the secure path stays buffered, mirroring the mmap read exclusion); every other backend/serializer falls back to set(bytes) unchanged; streamed values skip L1 by design. Measured (400MB incompressible frame, subprocess ru_maxrss): streaming write peaks at 2.29x logical vs 5.58x buffered. Locked in by a peak-RSS regression test alongside the #152/#171 harness. Closes #245
This comment has been minimized.
This comment has been minimized.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 55 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (5)
WalkthroughAdds a bounded-memory streaming write path for plaintext Arrow values. Arrow data is serialized directly to supported backend sinks, FileBackend commits temporary files atomically, handlers provide buffered fallback, and tests and documentation cover eligibility, limits, failures, and memory usage. ChangesStreaming Arrow write path
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CacheOperationHandler
participant CacheSerializationHandler
participant StandardCacheHandler
participant FileBackend
participant ArrowSerializer
CacheOperationHandler->>CacheSerializationHandler: check streaming eligibility
CacheOperationHandler->>StandardCacheHandler: set_streaming with payload callback
StandardCacheHandler->>FileBackend: stream payload to backend
FileBackend->>CacheSerializationHandler: invoke callback on temporary sink
CacheSerializationHandler->>ArrowSerializer: serialize_to_sink
ArrowSerializer->>FileBackend: write Arrow IPC bytes and patched checksum
FileBackend-->>CacheOperationHandler: commit result
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![kodus-27b[bot]](https://avatars.githubusercontent.com/u/223823690?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cachekit/backends/file/backend.py`:
- Around line 413-458: Refactor set_streaming so the producer callback and
temporary-file write run outside the process-wide self._lock, avoiding
serialization of other FileBackend operations during long streams. Reacquire
self._lock only around the final capacity check, atomic rename, and _maybe_evict
commit sequence, while retaining _acquire_file_lock for file-level safety and
cleaning up the temporary file on failures. Preserve and explicitly confirm the
intended capacity-check ordering now that it occurs after streaming.
- Around line 771-787: Update refresh_ttl to use the existing _expiry_from_ttl
helper for TTL validation and expiry calculation, removing its duplicated
zero-TTL handling, bounds checks, error construction, and time arithmetic while
preserving the helper’s established behavior.
In `@src/cachekit/cache_handler.py`:
- Around line 1490-1511: Update the Returns documentation for store_result at
src/cachekit/cache_handler.py lines 1490-1511 to distinguish bytes for a
buffered store, which is eligible for L1, from None for a streamed store that
intentionally skips L1 or for a failed store. Apply the same Returns
clarification to store_result_async at src/cachekit/cache_handler.py lines
1560-1578; do not change the store behavior.
In `@src/cachekit/serializers/arrow_serializer.py`:
- Around line 373-374: Add a concise comment immediately before the exception
handler around the Arrow IPC serialization path explaining that ValueError is
intentionally not wrapped because _HashingSink budget errors must propagate
unchanged for the handler and tests. Keep the existing ArrowInvalid and
ArrowTypeError wrapping behavior unchanged.
In `@tests/unit/test_streaming_write_path.py`:
- Around line 260-263: Update test_ttl_zero_and_none_mean_no_expiry to exercise
set_streaming with both ttl=0 and ttl=None, asserting each stored value remains
retrievable without expiry. Keep the test focused on verifying that both TTL
values mean no expiration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 1cd3aaf6-7b70-49bc-86bb-40d3182ef92d
📒 Files selected for processing (12)
.secrets.baselinedocs/backends/file.mddocs/configuration.mddocs/serializers/arrow.mdsrc/cachekit/backends/base.pysrc/cachekit/backends/file/backend.pysrc/cachekit/cache_handler.pysrc/cachekit/config/settings.pysrc/cachekit/serializers/arrow_serializer.pysrc/cachekit/serializers/wrapper.pytests/performance/test_large_object_memory.pytests/unit/test_streaming_write_path.py
- FileBackend.set_streaming no longer holds the process-wide lock for the whole stream: the producer runs the entire serialization (minutes at multi-GB scale), and safety comes from the unique temp path + flock + atomic rename, not the RLock. Lock now covers only the capacity precheck and the rename/eviction commit. - refresh_ttl reuses _expiry_from_ttl instead of duplicating TTL bounds. - store_result(_async) Returns docs: None is 'nothing for L1', not a failure signal (streamed success, streamed failure, or serialize error). - serialize_to_sink: comment why ValueError is deliberately not wrapped (mid-stream max_value_size budget contract). - test_ttl_zero_and_none_mean_no_expiry now actually covers ttl=0.
27Bslash6
changed the title
This comment has been minimized.
This comment has been minimized.
…p cleanup (LAB-766) - f: BinaryIO | None instead of Any (type safety preserved on the handle). - Replace the BaseException catch-and-reraise with a committed-flag finally: same guarantee (no temp-file leak on any exit path, incl. KeyboardInterrupt), no overly-broad catch, producer exceptions still propagate unwrapped.
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@kody start-review |
Kody Review CompleteGreat news! 🎉 Keep up the excellent work! 🚀 Kody Guide: Usage and ConfigurationInteracting with Kody
Current Kody ConfigurationReview OptionsThe following review options are enabled or disabled:
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
…B-766) Textual conflict in tests/performance/test_large_object_memory.py: kept both the new streaming-write RSS test (this branch) and the skipif decorator LAB-350 added to test_byte_storage_store_has_no_full_payload_copy. Silent (non-textual) conflict in arrow_serializer.py: LAB-110 replaced _bounded_chunksize with _write_bounded_batches on main to fix skewed-frame batch overshoot (cachekit-py#161), but this branch's serialize_to_sink still called the removed function. Switched it to _write_bounded_batches to match serialize() and keep the byte-aware batching fix on the streaming path too.
a85c9fa
1 participant
Closes #245. Tracked as LAB-766.
What
The write path materialized the whole serialized envelope before a single byte reached the backend — pyarrow's
BufferOutputStreamdoubling-growth + memory-pool retention held ~3.75x the payload, putting write peak RSS at ~5.6x logical for large DataFrames (the residual PR #152 left behind). This is the write-side twin of theget_buffermmap read path (#171 / #187): stream serialized record batches straight to the backend, never holding the full blob.Measured (400MB incompressible frame, subprocess
ru_maxrss): streaming write peaks at 2.29x logical vs 5.58x buffered. Locked in by a peak-RSS regression test (bound 3.2x — between the two with margin on both sides).How
BufferWritableBackendcapability protocol (backends/base.py), same@runtime_checkable+ TypeGuard idiom asBufferReadableBackend. Writer-callback shape —set_streaming(key, write_payload, ttl)against a seekable sink — rather than anIterable[bytes], because the Arrow envelope stores its 8-byte xxHash3-64 checksum BEFORE the IPC bytes it covers: a forward-only chunk iterator cannot express that without serializing twice. The producer hashes incrementally while streaming and patches the checksum in place, so the stored envelope is byte-identical to the buffered path (asserted in tests) and every existing read path (get, mmapget_buffer) works unchanged.FileBackend.set_streaming: incremental write behind the existing temp-file + exclusive-flock + fsync + atomic-rename pattern. Producer failures discard the partial write, leave any previous value for the key intact, and propagate unwrapped.max_value_mbenforced on the on-disk size before rename.ArrowSerializer.serialize_to_sink: streams the ~8 MiB record-batch IPC frames through a hashing tee into the sink, bypassingBufferOutputStreamentirely. Optional byte budget aborts over-limit values mid-stream (the issue L1 budget is not configurable + three dead/unenforced config knobs #163max_value_sizeceiling, enforced incrementally instead of post-materialization).store_result(_async): streams only for plaintext Arrow on a streaming-capable backend. Everything else — Redis, CachekitIO, L1, memcached, non-Arrow serializers, SWR writes — keeps usingset(bytes)with no behaviour change (byte-for-byte, asserted).Not a protocol/crypto change
SerializationWrapperis Python-SDK-internal and opaque to other SDKs and the SaaS; the cross-SDK ByteStorage wire format, AAD construction, key derivation, and cache-key format are untouched.wrap()was refactored towrap_prefix() + payloadwith byte-identical output (covered by existing wrapper doctests + new equality tests).Docs pass
docs/backends/file.md: new "Bounded-Memory Large Values (Arrow)" section (both capabilities, measured numbers, encryption carve-out, L1 exclusion).docs/serializers/arrow.md: write-side streaming added to the Memory Usage section.docs/configuration.md+settings.py(arrow_compression): clarified that write streaming is codec-independent; mmap reads still requirenone.backends/base.pymodule docstring: capability-protocol list brought up to date.Testing
tests/unit/test_streaming_write_path.py): eligibility, delegation, byte-identity vs buffered, fallback, no-buffered-retry-on-failure, atomicity (old value survives failed stream, no temp litter), TTL bounds,max_value_mb, mmap-read-of-streamed-entry composition, async twins.tests/performance/test_large_object_memory.py), same subprocess harness as the fix: bound memory for large DataFrame/Arrow caching (was OOMing at real sizes) #152/EPIC: Zero-copy mmap read path for the File backend (large-object read RSS) #171 tests.ruff check/ruff format/ basedpyright (pre-commit) pass. Remaining failures intests/integration/saas+ timing-sensitive perf benchmarks reproduce on the clean base (environmental).Summary by CodeRabbit
New Features
Documentation