perf: defer per-super-step read_channels deep clone in non-streaming mode - #31
Open
MuFengMuXue wants to merge 20 commits into
Open
perf: defer per-super-step read_channels deep clone in non-streaming mode#31MuFengMuXue wants to merge 20 commits into
MuFengMuXue wants to merge 20 commits into
Conversation
The messages reducer is the default for the `messages` channel in every agent state. Each LLM turn previously deep-cloned the entire accumulated history (all message content + tool args) — O(n^2) over a run. Change the reducer signature from `fn(&JsonValue, &JsonValue) -> JsonValue` to `fn(JsonValue, &JsonValue) -> JsonValue` and have BinaryOperatorAggregate::update hand ownership of the accumulated value to the reducer via `guard.take()`, so existing messages are moved instead of cloned. `update` stays borrowed (usually a single new message). add_messages now: - consumes `current` and moves existing messages into the result - collects remove-ids as borrowed `&str` (no per-message String alloc) - takes a fast path when there are no removals A benchmark is added (tests/bench_messages.rs, #[ignore]d) measuring the reducer hot path: ~100x faster at 250 steps up to ~330x at 1000 steps, and scaling drops from quadratic to near-linear. This is a breaking API change for reducers: they now receive `current` by value.
Checkpoint saves re-encoded and re-inserted a blob row for every channel on every super-step. With a large static channel (e.g. embedded context) that meant re-serializing and re-writing the same value each step — pure write amplification. save_checkpoint now takes the channel versions as of the start of the run and derives new_versions = channels whose version moved, passing only the delta to the saver. SqliteSaver::put writes blob rows only for new_versions; reads merge version-joined blob values over the checkpoint row body (extend, not replace) so every channel resolves through the version join regardless of which checkpoint wrote the blob. Adds test_incremental_blob_writes (a second checkpoint with one updated channel writes one new blob row; both new and old checkpoints read back fully) and a benchmark suite (tests/bench_pregel.rs, #[ignore]d) covering checkpointed growth and a static-context scenario. Measured (release, 800 steps, per-step): sqlite linear 12.651->12.230ms (~3%), static 200KB context 12.723->12.200ms (~4%). The win is small on these micro-benches because the dominant per-step cost is in-process state handling (read_channels deep clone), not blob I/O; write amplification is eliminated, which matters far more for remote savers and large static channels.
MuFengMuXue
added a commit
to MuFengMuXue/langgraph-rust
that referenced
this pull request
Aug 2, 2026
…asurement - LatestOnlySaver retains only the newest checkpoint per thread. InMemorySaver keeps every checkpoint, so the growth benches retained O(steps²) serialized state and OOM'd after a few hundred steps; retention is now O(latest state). - The multi-step loop's conditional routing only sees the node output (the combined PregelNode evaluates branches on the node delta, not the accumulated state), so routing on messages.len() always saw a single-element array and the loop ran to the recursion limit — silently truncating every run at 25 super-steps (the default limit) and invalidating the earlier perf claim. Route on an explicit count channel so the loop runs the real target super-steps, and the final sanity assert catches any truncation. - Re-measured Onelevenvy#31 on the corrected bench: ~10-22% faster at 50/100/200 real super-steps; the absolute saving grows with history.
MuFengMuXue
added a commit
to MuFengMuXue/langgraph-rust
that referenced
this pull request
Aug 2, 2026
…asurement - LatestOnlySaver retains only the newest checkpoint per thread. InMemorySaver keeps every checkpoint, so the growth benches retained O(steps²) serialized state and OOM'd after a few hundred steps; retention is now O(latest state). - The multi-step loop's conditional routing only sees the node output (the combined PregelNode evaluates branches on the node delta, not the accumulated state), so routing on messages.len() always saw a single-element array and the loop ran to the recursion limit — silently truncating every run at 25 super-steps (the default limit) and invalidating the earlier perf claim. Route on an explicit count channel so the loop runs the real target super-steps, and the final sanity assert catches any truncation. - Re-measured Onelevenvy#31 on the corrected bench: ~10-22% faster at 50/100/200 real super-steps; the absolute saving grows with history.
MuFengMuXue
force-pushed
the
feat/lazy-read-channels
branch
from
August 2, 2026 10:22
f0e6dc2 to
16d792b
Compare
…asurement - LatestOnlySaver retains only the newest checkpoint per thread. InMemorySaver keeps every checkpoint, so the growth benches retained O(steps²) serialized state and OOM'd after a few hundred steps; retention is now O(latest state). - The multi-step loop's conditional routing only sees the node output (the combined PregelNode evaluates branches on the node delta, not the accumulated state), so routing on messages.len() always saw a single-element array and the loop ran to the recursion limit — silently truncating every run at 25 super-steps (the default limit) and invalidating the earlier perf claim. Route on an explicit count channel so the loop runs the real target super-steps, and the final sanity assert catches any truncation. - Re-measured Onelevenvy#31 on the corrected bench: ~10-22% faster at 50/100/200 real super-steps; the absolute saving grows with history.
MuFengMuXue
force-pushed
the
feat/lazy-read-channels
branch
from
August 2, 2026 10:26
16d792b to
6006a43
Compare
MuFengMuXue
added a commit
to MuFengMuXue/langgraph-rust
that referenced
this pull request
Aug 2, 2026
…he multi-step loop measurement - LatestOnlySaver retains only the newest checkpoint per thread. InMemorySaver keeps every checkpoint, so the growth benches retained O(steps²) serialized state and OOM'd after a few hundred steps; retention is now O(latest state). - The multi-step loop's conditional routing only sees the node output, so routing on messages.len() always saw a single-element array and the loop ran to the recursion limit (25 default) — silently truncating every run. Route on an explicit count channel so the loop runs the real target super-steps. - Same bench content as PR Onelevenvy#31's correction commit (6006a43) so either PR merging first lands the identical corrected bench.
Thread the channel set returned by apply_writes into prepare_next_tasks as the candidate set, so super-steps after the first only scan nodes reachable from bumped channels instead of all N nodes. The first super-step keeps a full scan because START-edge and input writes bypass apply_writes. Linear chain of 1000 no-op nodes: 116.8ms -> 46.9ms (~2.5x). Still super-linear because apply_writes itself sweeps all channels each super-step (step-2 max_by over channel_versions, step-6 update([]) notify sweep). Correctness: bench_sequential_chain asserts done==true at 1000 nodes (no truncation); 37 langgraph-core-rs tests pass.
Three changes on the chain hot path, together taking a 1000-node no-op chain from 46.9ms -> 4.28ms (~11x; vs the 116.8ms baseline ~27x, now within ~1.8x of juncture): 1. apply_writes takes next_version from the caller instead of scanning every channel version for the max each superstep. run_pregel_inner keeps a running version counter (initialized once; versions only ever grow). 2. The step-6 notify sweep touches only the previous super-step's 'updated' channels instead of every channel. A channel can only be available at step-6 if it was written last superstep (that sweep clears everything else), so the bounded sweep is behaviorally identical. The first superstep and get_state pass None (full sweep). 3. version_gt fast path: equal-length all-digit strings compare lexically, skipping the f64 parse and fixing precision loss past 2^53 for the 32-digit engine versions. get_state's call site keeps its exact old behavior. Correctness: all workspace tests pass except the pre-existing sqlite failure; bench_sequential_chain asserts done==true at 1000 nodes (no truncation); fanout/multi-step/sqlite benches unchanged.
build_pregel_nodes + build_trigger_to_nodes are pure functions of the immutable CompiledStateGraph fields, yet were rebuilt on every invoke. Measured ~1.4ms at 1000 no-op nodes (~1/3 of a chain's runtime). Build both once in compile_with() and store on the struct; run_pregel_inner, get_state and get_state_history use the cached refs. PregelNode gains #[derive(Clone)] (all Arc/Vec/String fields). 1000-node no-op chain: 4.32ms -> 2.56ms (~1.7x), ~2.56us/node = juncture parity.
Fan-out super-steps now take ~max(branch time) instead of the sum: - single task executes inline (no spawn overhead on the chain path) - >=2 tasks dispatched through tokio::task::JoinSet - tasks re-sorted to index order on return so StreamMode::Updates emission and apply_writes channel ordering stay deterministic - lowest task-index error/interrupt wins regardless of JoinSet completion order; panics surface as a generic TaskFailed - run_tasks now returns (tasks, result) so the caller keeps completed writes on error/interrupt (interrupt path unchanged) Bench: 2/4/8 branches x 40ms = 54/48/46ms (was 93/185/371ms serial). Chain path unaffected (single-task inline).
…ridge - BaseCheckpointSaver::put takes Checkpoint by value (moves instead of deep-copying on the hot per-superstep save path) - InMemorySaver stores the Checkpoint struct directly — no to_value/ from_value round-trip per step; per-thread latest-pointer makes get/put latest lookups O(1) instead of scanning the whole history - async trait defaults only block_in_place on multi-thread runtimes (fixes current_thread panic); in-memory savers get true async overrides - run_pregel_inner/get_state/get_state_history/update_state call aput/aget_tuple/aput_writes, so the per-superstep save never bridges through block_in_place - channels_from_checkpoint consumes the checkpoint's channel values instead of cloning
build_checkpoint now serializes only the channels whose version moved since the run started (new_versions delta) instead of deep-copying every channel's checkpoint() each super-step. Unchanged channels are reconstructed by the saver's version-merged reads, so nothing is lost. - InMemorySaver: store channel values as version-addressed blobs (rows + blobs + O(1) latest), reconstruct the full channel_values on get_tuple/list via reconstruct_values. Mirrors the sqlite blob model. - LatestOnlySaver (bench): merge the delta into the retained checkpoint on put instead of replacing it wholesale. - run_pregel_inner input handling: bump versions for every channel written from the input (both START-mapped and direct dict-key writes). Previously dict-key writes set the value without a version entry, so a value-bearing channel could be silently dropped from a delta-only checkpoint (regression surfaced by test_update_state).
Cache maps (thread_id, checkpoint_ns, channel, version_str) -> JsonValue for savers that store channel values as version-addressed blobs. Repeated reads of an unchanged large channel (e.g. a static context) become O(clone) instead of O(DB fetch + parse). Correctness relies on version strings being monotonic within a (thread_id, checkpoint_ns) pair; the sole reuse point is delete_thread + thread recreation, so callers MUST call remove_thread there. Whole-cache clear on overflow bounds memory. Not yet wired into any saver (wired into the sqlite read path next).
load_blobs now drives reads by the checkpoint's channel_versions, serving cache hits (immutable per (channel, version)) as clones and fetching only the misses from the DB via a (channel, version) IN query. Repeated reads of an unchanged large channel (e.g. a static context) become O(clone) instead of O(DB fetch + parse) — the remaining P0 read-side bottleneck. - SqliteSaver gains an Arc<BlobCache> field (shared across clones). - aget_tuple/alist pass channel_versions through; alist's rows of one thread share stable versions, so all but the first are cache hits. - adelete_thread calls BlobCache::remove_thread (required: versions restart at 1 for a recreated thread with the same id). - version_to_string helper shared by dump_blobs and load_blobs keeps write keys and read lookups in sync; SELECT_BLOBS_SQL dropped. - Cache is populated only from the read path (never warmed from aput), so cached values equal exactly what the old parse produced. - Tests: cache hit served with no blob rows in the DB; delete_thread invalidation; empty channel_versions; cleared-channel not returned.
- bench_sqlite_static_context asserts the 128KB static context survives every step and messages accumulate to 2*steps (each invoke appends the input message plus the node's append, matching sanity_bench_graphs_work). - bench_sqlite_static_context_reads: seed once, time one cold get_state (cache miss, parses the 128KB blob) vs N warm reads (cache hits, clone). Measured ~558us cold vs ~72us warm = ~7.7x read-side speedup from the BlobCache.
…ints The sync invoke() entry points (CompiledStateGraph::invoke, RunnableCallable::invoke) built and dropped a full multi-thread Runtime on every call from outside a tokio context — each construction spawns num_cpus worker threads (~100us+). A single OnceLock-cached multi-thread runtime amortizes that; Handle::try_current() guard is preserved, so callers already inside a runtime keep using their own handle (no create/cache). Also unifies RunnableCallable::invoke's cold path to install the task-local config via with_config (matching ainvoke), fixing a pre-existing inconsistency where the cold path skipped it. New bench_sync_invoke_runtime_cache probe (plain #[test], outside tokio): 820.7us/invoke -> 3.3us/invoke (~248x) on the sync invoke path.
…fig in place execute_task / execute_single_task_sync cloned task.config to inject CONFIG_KEY_SEND, leaving the original untouched. task.config is dead after the super-step (nothing downstream reads it), so mutate it in place: the no-store/no-stream hot path now makes 0 config clones per task, and the with_runtime branch drops from 2 clones to 1 (the remaining clone is structural — with_runtime owns a config for the task-local while the node borrows one). bench_sequential_chain @1000 nodes: 2.5909ms -> 2.3974ms (2.59 -> 2.397 us/node, ~7.5%). bench_react_agent: 50-turn ~3% better, 100/200 within noise (history serialization dominates).
apply_writes grouped writes by channel via `val.clone()` per write while borrowing tasks immutably. Since the streaming Updates emit (state.rs) runs over task.writes before apply_writes, and the get_state snapshot path now computes `next` (writes.is_empty()) before the call, nothing reads task.writes afterward — so take each task's write buffer by value (&mut [PregelExecutableTask] + mem::take) and move the values instead of cloning them. get_state's `next` computation is reordered before the call; the interrupt-only apply_completed_writes path is untouched (cold, mutually exclusive). Correctness: 37 core + 14 prebuilt tests green. Measurement: chain and multi-step-loop benches are within run-to-run noise on small writes (the saving is one value-clone per write, which grows with write size — worth the most for large node outputs).
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.
Summary
Defers the per-super-step
read_channelsdeep clone inrun_pregel_inner. In non-streaming mode the output was re-read (deep-cloning every output channel's value) every super-step — the dominant per-step cost at large message histories. It is now deferred to the two points that actually consume it: theinterrupt_afterreturn and the loop exit. In streaming mode the already-emittedstream_mode="values"value is reused instead of re-reading.ran_super_steppreserves the originalNullresult for empty runs (empty graph / fork with no resume).Benchmark
tests/bench_pregel.rs(#[ignore]d) —bench_multi_step_loopexercises the affected path directly.The bench was corrected in
6006a43before these numbers were taken:messages.len()from the node output, which is always a single-element array (the combinedPregelNodeevaluates conditional branches on the node delta, not the accumulated state) — so every run silently truncated at the recursion limit (25 by default) instead of running the requested super-steps. It now routes on an explicitcountchannel and runs the real 50/100/200 super-steps; the sanity assert catches any truncation.LatestOnlySaver(retains only the newest checkpoint per thread) replacesInMemorySaver, which kept every checkpoint and caused O(steps²) retention / OOM in the growth benches.Measured on the corrected bench (Windows, mimalloc, same machine, pre- vs post-commit
state.rs):The absolute saving grows with history length (~0.8 ms at 50 → ~9 ms at 200); run-to-run noise is ~±10% on this machine.
Validation
cargo fmt --all -- --check✅cargo clippy --all-targets --all-features -- -D warnings✅cargo test --workspace✅ (one pre-existing unrelated failure inlanggraph-checkpoint-sqlite, reproduced on clean main)cargo test --release --test bench_pregel -- --ignored✅ (all benches, memory bounded to 1 retained checkpoint per thread)