Memory Aware Caching in Lighthouse #7803

PoulavBhowmick03 · 2025-07-27T18:29:48Z

WIP

Issue Addressed

Fixes #7449

Changes

Memory tracking

Introduced BeaconStateMemorySize (in store/memsize.rs) that wraps a BeaconState in BeaconStateWrapper and measures with Milhouse MemoryTracker.
Added a histogram metric beacon_state_memory_size_calculation_time to time memory sizing calls.

State cache API/behavior

StateCache now has:
- max_cached_bytes: configurable byte ceiling for the hot state cache.
- cached_bytes: last measured total (not incrementally maintained).
- put_count: counter to throttle re-measurement.
- needs_remeasure: boolean, to check if there is need for remeasuring of the memory_size
put_state:
- Every RECOMPUTE_INTERVAL inserts calls the full re-measurement routine, if needs_remeasure is true.
delete_state and related paths:
- Remove entries and mappings only (no byte subtraction).. size is enforced by periodic measurement of the entire cache.
Added measure_cached_memory_size():
- Runs a single MemoryTracker across all cached states and sums differential_size to avoid double-counting shared Arcs.
Added recompute_cached_bytes():
- Measures if over max_cached_bytes, prune in batches (uses headroom, clamped to 5–64) → re-measure, loop until under the limit.
- Updates gauge store_beacon_state_cache_memory_size and logs measurement time.

Config / CLI / metrics

New CLI flag: --state-cache-max-bytes (default 512 MiB).
New constant: DEFAULT_STATE_CACHE_MAX_BYTES.
HotColdDB updates the memory gauge with state_cache.cached_bytes().
Tests updated to assert the default for max_state_cache_bytes.

dapplion · 2025-08-04T17:15:56Z

beacon_node/store/src/state_cache.rs

            self.states.put(state_root, (state_root, state.clone()))
        {
            deleted_states.push(deleted_state_root);
+            self.cached_bytes = self
+                .cached_bytes
+                .saturating_sub(removed_state.memory_size());


@michaelsproul can comment but can you just count the size of a single state? Does memory_size only count the size of the data owned by the first copy of the Arc?

Can you add timer metrics on memory_size to measure how expensive it is in terms of time?

Will do so. And the memory_size implementing Milhouse's MemoryTracker only counts the data owned by Arc, once, if that's what you meant? It does not overcount the shared Arc data

Can you add timer metrics on memory_size to measure how expensive it is in terms of time?

The memory_size of a state is probably larger than the amount of space we save by deleting it, because some of its memory might be shared with other states.

I don't think it's feasible to store enough information to update the cached_bytes after removing a state – we need to just recompute the cache size. This is one of the challenges of this approach. You fundamentally don't know how much memory you will recover unless you recount.

Maybe I'm wrong though? Maybe we could change the order that we count states so that the ones most likely to be pruned are tracked last. Then a state's differential size would be approximately how much memory we'd be able to recover by deleting it.

…memory

… iteratively track memory usage

…acking

michaelsproul · 2025-09-01T05:51:13Z

beacon_node/beacon_chain/src/chain_config.rs

@@ -10,7 +10,7 @@ pub const DEFAULT_RE_ORG_MAX_EPOCHS_SINCE_FINALIZATION: Epoch = Epoch::new(2);
 /// Default to 1/12th of the slot, which is 1 second on mainnet.
 pub const DEFAULT_RE_ORG_CUTOFF_DENOMINATOR: u32 = 12;
 pub const DEFAULT_FORK_CHOICE_BEFORE_PROPOSAL_TIMEOUT: u64 = 250;
-
+pub const DEFAULT_STATE_CACHE_MAX_BYTES: usize = 536870912;


Just a stylistic thing but we could write this constant as 512 * (1 << 20), perhaps with an intermediate constant like const MEGABYTE: usize = 1 << 20.

Or 512 * 1024 * 1024 is even better!

This config can probably be deleted actually, as you've got the default in the store config

michaelsproul · 2025-09-01T05:52:01Z

beacon_node/src/cli.rs

+                .long("state-cache-max-bytes")
+                .value_name("BYTES")
+                .help("Specifies the maximum size of the state cache in bytes")
+                .default_value("536870912")


This CLI param should probably be in MiB, as I doubt it's useful to set a finer granularity than 1 MiB.

michaelsproul · 2025-09-01T05:52:25Z

beacon_node/src/config.rs

@@ -377,6 +377,12 @@ pub fn get_config<E: EthSpec>(
            .map_err(|_| "state-cache-size is not a valid integer".to_string())?;
    }

+    if let Some(max_bytes) = cli_args.get_one::<String>("state-cache-max-bytes") {


parse_optional will be slightly less verbose here

michaelsproul · 2025-09-01T05:56:54Z

beacon_node/store/src/memsize.rs

+    }
+
+    fn subtrees(&self) -> Vec<&dyn MemorySize> {
+        vec![]


You're missing all of the fields of the BeaconState, i.e. the interesting parts. See my previous draft PR:

https://github.com/michaelsproul/lighthouse/blob/e34fafab5977faae5362c3cc16f969256769502e/consensus/types/src/beacon_state/memsize.rs#L17-L90

michaelsproul · 2025-09-01T05:57:57Z

beacon_node/store/src/memsize.rs

+    fn memory_size(&self) -> usize {
+        let wrapper = BeaconStateWrapper(self);
+        // Timer for MemorySize
+        let timer = Instant::now();


You can do this using metrics::start_timer and drop (no need for manual Instant use)

michaelsproul · 2025-09-01T06:01:02Z

beacon_node/store/src/state_cache.rs

+            self.cached_bytes = self.cached_bytes.saturating_sub(
+                self.states
+                    .peek(state_root)
+                    .map_or(0, |(_, state)| state.memory_size()),


I don't think we can measure the size of a removed state like this, because memory_size as written will just calculate the size of this state in isolation, which ignores the structural sharing.

I think it's probably better just to not even try to update this on removal, because it would require a recompute. We can just wait for the next recompute (or trigger a recompute after deletion).

Yeah, looks like you're already doing multiple recalcs inside cull, so this can just be deleted

michaelsproul · 2025-09-01T06:04:13Z

beacon_node/store/src/state_cache.rs

+            }
+
+            // Recalculate the memory size after culling.
+            let (new_total_bytes, elapsed_time) = self.measure_cached_memory_size();


Doing multiple recalcs in a loop might be a bit slow, but we can assess once MemorySize is implemented for real

i.e. I expect the calulcation to take a lot longer once we actually examine the fields of the BeaconState, this issue:

#7803 (comment)

PoulavBhowmick03 requested a review from michaelsproul as a code owner July 27, 2025 18:29

PoulavBhowmick03 marked this pull request as draft July 27, 2025 18:31

PoulavBhowmick03 mentioned this pull request Jul 27, 2025

Measure state cache size in memory #7449

Open

dapplion reviewed Aug 4, 2025

View reviewed changes

michaelsproul assigned dapplion Aug 11, 2025

PoulavBhowmick03 force-pushed the memory_aware_caching branch 2 times, most recently from fd74766 to 257938d Compare August 21, 2025 09:29

PoulavBhowmick03 added 10 commits August 29, 2025 11:03

Fixed merge conflicts

07a19bf

memory tracking in state_cache.rs and modified put_state fn to track …

b2aeb85

…memory

register_metrics updated to add STORE_BEACON_STATE_CACHE_MEMORY_SIZE

5074f3a

modified delete function and added cli argument

e65d834

removed incremental changes to cached_bytes and added new function to…

13629c0

… iteratively track memory usage

Fixed merge conflicts and added a new fn to avoid share arc memory tr…

d681be3

…acking

Fixed conflicts

b6874b8

fixes CI

be7a772

fixes CI

cf9ef9b

wrote benchmarks for state_cache and ran to check for regressions

d758cdb

PoulavBhowmick03 force-pushed the memory_aware_caching branch from 05166bd to d758cdb Compare August 31, 2025 08:41

PoulavBhowmick03 marked this pull request as ready for review August 31, 2025 10:49

michaelsproul reviewed Sep 1, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Memory Aware Caching in Lighthouse #7803

Memory Aware Caching in Lighthouse #7803

Uh oh!

PoulavBhowmick03 commented Jul 27, 2025 •

edited

Loading

Uh oh!

dapplion Aug 4, 2025

Uh oh!

dapplion Aug 4, 2025

Uh oh!

PoulavBhowmick03 Aug 4, 2025 •

edited

Loading

Uh oh!

michaelsproul Aug 5, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

michaelsproul Sep 1, 2025

Uh oh!

Uh oh!

Memory Aware Caching in Lighthouse #7803

Are you sure you want to change the base?

Memory Aware Caching in Lighthouse #7803

Uh oh!

Conversation

PoulavBhowmick03 commented Jul 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Issue Addressed

Changes

Memory tracking

State cache API/behavior

Config / CLI / metrics

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

PoulavBhowmick03 Aug 4, 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!

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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

PoulavBhowmick03 commented Jul 27, 2025 •

edited

Loading

PoulavBhowmick03 Aug 4, 2025 •

edited

Loading