From f959cefb6a2ddf53489ae93f68fec673abc66761 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 29 Jul 2026 18:03:56 -0400 Subject: [PATCH] fix(model cache): release shared weights when a cache goes away MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing released a cache's SharedCpuWeightsStore references except _delete_cache_entry(): shutdown() left every resident record's refcount held, and a cache dropped without shutdown() (test teardown; any future wiring that rebuilds caches at runtime) stranded the canonical tensors and their accounting forever. Today's production wiring tears the store down together with its caches, so the live exposure is cross-test pollution of the process-global store and RAM pinned past ModelManagerService.stop() — but the refcount invariant ('every acquire is paired with exactly one release') was simply not upheld, and this makes it self-healing before any wiring change turns it into a real peer-accounting bug. Two mechanisms, for the two ways a cache goes away: - shutdown() now releases its resident records' shared references synchronously — it runs in a normal thread context, so the direct (locking) release is safe there, and teardown does not depend on a later store operation happening. - Each wrapper registers a weakref.finalize fallback for the dropped-without-shutdown case. The finalizer runs in GC context, where taking the store's non-reentrant lock could self-deadlock (a collection can fire inside acquire()'s critical section on the same thread — the rule ModelCache.release_first_use_grace documents), so it only ENQUEUES into a SimpleQueue; every public store method drains the queue under the lock. The finalizer is registered inside the acquire's try (a registration failure must release too), its args carry the key and canonical dict rather than the wrapper (finalize holds args strongly — referencing self would make the wrapper immortal), and release_shared_weights() detaches it before releasing synchronously so eviction-then-collection releases exactly once. The state-dict identity keeps releases correct across invalidate()'s retired entries. RamBudget.total_in_use() now documents why its store read must stay outside the budget lock: the drain allocates under the store lock, so GC can run _on_cache_collected (store→budget) there, and a budget→store order anywhere would complete the deadlock cycle. Six regression tests, verified to fail before the fix, covering: shutdown releases synchronously with an empty queue; collection returns refcount/bytes/budget to zero; the collection-time release is enqueue-only (never applied inline by GC); eviction + collection release exactly once across two caches; a retired (invalidated) entry is freed by a collected holder; and the partial-load wrapper behaves like the full-load one. One existing test relied on an abandoned wrapper leaking its reference and now binds it. Co-Authored-By: Claude Fable 5 --- .../cached_model_only_full_load.py | 26 +++- .../cached_model_with_partial_load.py | 26 +++- .../load/model_cache/model_cache.py | 7 + .../load/model_cache/ram_budget.py | 5 + .../load/model_cache/shared_cpu_weights.py | 73 +++++++-- .../test_cached_model_shared_weights.py | 6 +- .../test_model_cache_ram_budget.py | 145 ++++++++++++++++++ 7 files changed, 268 insertions(+), 20 deletions(-) diff --git a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py index ccf40654575..226f65f6998 100644 --- a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py +++ b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py @@ -1,3 +1,4 @@ +import weakref from typing import Any import torch @@ -41,6 +42,7 @@ def __init__( # under `cache_key`; `release_shared_weights()` must be called exactly once on eviction. self._shared_store: SharedCpuWeightsStore | None = None self._shared_key: str | None = None + self._shared_release_finalizer: weakref.finalize | None = None # A CPU read-only copy of the model's state dict. self._cpu_state_dict: dict[str, torch.Tensor] | None = None @@ -57,9 +59,24 @@ def __init__( if canonical is not cpu_state_dict: model.load_state_dict(canonical, assign=True) cpu_state_dict = canonical + # A cache dropped without a shutdown() never routes its records through + # _delete_cache_entry, so nothing would call release_shared_weights() and the + # canonical tensors would stay resident (and counted by the RAM budget) + # forever. The finalizer must not reference `self` (its args are held strongly + # — that would make the wrapper immortal) and must not take the store's + # non-reentrant lock (it runs in GC context): release_deferred only enqueues; + # the store applies it on its next operation. release_shared_weights() detaches + # this on the normal eviction path, so the release happens exactly once either + # way. Registered inside this try so a failure here (e.g. MemoryError) + # releases the just-acquired reference too. + self._shared_release_finalizer = weakref.finalize( + self, shared_store.release_deferred, cache_key, canonical + ) + self._shared_release_finalizer.atexit = False except Exception: - # The re-point failed after acquiring a reference; release it so the shared - # entry's refcount isn't leaked (this wrapper will never enter the cache). + # The re-point or finalizer registration failed after acquiring a reference; + # release it so the shared entry's refcount isn't leaked (this wrapper will + # never enter the cache). self.release_shared_weights() raise self._cpu_state_dict = cpu_state_dict @@ -92,6 +109,11 @@ def release_shared_weights(self) -> None: no-op. After release, the shared store frees the canonical tensors once the last device that held this key releases it. """ + if self._shared_release_finalizer is not None: + # The eviction path is releasing synchronously; the collection-time fallback must not + # release the same reference a second time. + self._shared_release_finalizer.detach() + self._shared_release_finalizer = None if self._shared_store is not None and self._shared_key is not None: self._shared_store.release(self._shared_key, self._cpu_state_dict) self._shared_store = None diff --git a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py index cae68e24331..00be7d3d950 100644 --- a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py +++ b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py @@ -1,3 +1,5 @@ +import weakref + import torch from invokeai.backend.model_manager.load.model_cache.shared_cpu_weights import SharedCpuWeightsStore @@ -29,6 +31,7 @@ def __init__( # under `cache_key`; `release_shared_weights()` must be called exactly once on eviction. self._shared_store: SharedCpuWeightsStore | None = None self._shared_key: str | None = None + self._shared_release_finalizer: weakref.finalize | None = None # Assigned for real at the end of __init__; initialized here so the acquire-failure path # below can call release_shared_weights(), which reads it, before that assignment runs. self._cpu_state_dict: dict[str, torch.Tensor] | None = None @@ -69,9 +72,23 @@ def __init__( if canonical is not cpu_state_dict: self._model.load_state_dict(canonical, assign=True) cpu_state_dict = canonical + # A cache dropped without a shutdown() never routes its records through + # _delete_cache_entry, so nothing would call release_shared_weights() and the + # canonical tensors would stay resident (and counted by the RAM budget) forever. + # The finalizer must not reference `self` (its args are held strongly — that would + # make the wrapper immortal) and must not take the store's non-reentrant lock (it + # runs in GC context): release_deferred only enqueues; the store applies it on its + # next operation. release_shared_weights() detaches this on the normal eviction + # path, so the release happens exactly once either way. Registered inside this try + # so a failure here (e.g. MemoryError) releases the just-acquired reference too. + self._shared_release_finalizer = weakref.finalize( + self, shared_store.release_deferred, cache_key, canonical + ) + self._shared_release_finalizer.atexit = False except Exception: - # The re-point failed after acquiring a reference; release it so the shared entry's - # refcount isn't leaked (this wrapper will never be inserted into the cache). + # The re-point or finalizer registration failed after acquiring a reference; + # release it so the shared entry's refcount isn't leaked (this wrapper will never + # be inserted into the cache). self.release_shared_weights() raise @@ -175,6 +192,11 @@ def release_shared_weights(self) -> None: no-op. After release, the shared store frees the canonical tensors once the last device that held this key releases it. """ + if self._shared_release_finalizer is not None: + # The eviction path is releasing synchronously; the collection-time fallback must not + # release the same reference a second time. + self._shared_release_finalizer.detach() + self._shared_release_finalizer = None if self._shared_store is not None and self._shared_key is not None: self._shared_store.release(self._shared_key, self._cpu_state_dict) self._shared_store = None diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 3950ed22eab..62be14cc023 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -542,6 +542,13 @@ def shutdown(self) -> None: if self._timeout_timer is not None: self._timeout_timer.cancel() self._timeout_timer = None + # Release the resident records' shared-weights references now, synchronously. A shut-down + # cache serves no more loads, and waiting for collection would leave the store's refcounts + # (and canonical tensors) to the wrappers' finalizers — which only ENQUEUE, and at teardown + # there may be no later store operation to drain the queue. shutdown() runs in a normal + # thread context, so the direct (locking) release is safe here. + for cache_entry in self._cached_models.values(): + cache_entry.cached_model.release_shared_weights() @synchronized @record_activity diff --git a/invokeai/backend/model_manager/load/model_cache/ram_budget.py b/invokeai/backend/model_manager/load/model_cache/ram_budget.py index fdc194ebe2a..55b3c3d2603 100644 --- a/invokeai/backend/model_manager/load/model_cache/ram_budget.py +++ b/invokeai/backend/model_manager/load/model_cache/ram_budget.py @@ -111,6 +111,11 @@ def remove_non_shared(self, nbytes: int, cache: Optional["ModelCache"] = None) - def total_in_use(self) -> int: """The true total RAM used by the model caches: shared weights (counted once) + non-shared.""" + # The store read MUST stay outside self._lock. The store's deferred-release drain runs + # under the store lock and allocates, so a cyclic GC can fire there and run + # _on_cache_collected, which takes THIS lock (store → budget on one thread). If any thread + # held the budget lock while calling into the store (budget → store), the two orders would + # deadlock against each other. shared = self._store.total_bytes_in_use() if self._store is not None else 0 with self._lock: non_shared = self._non_shared_bytes diff --git a/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py b/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py index 3e1fc9ad512..2992e1adfb5 100644 --- a/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py +++ b/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py @@ -1,3 +1,4 @@ +import queue import threading from dataclasses import dataclass, field @@ -50,6 +51,14 @@ class SharedCpuWeightsStore: def __init__(self) -> None: self._lock = threading.Lock() + # Releases posted from GC context (a cached-model wrapper's weakref.finalize when its cache + # was dropped without shutdown()/clear()). A finalizer must not take self._lock: it can run + # at any allocation point on any thread — including inside acquire()'s critical section, + # where taking this non-reentrant lock again would self-deadlock the process (see + # ModelCache.release_first_use_grace for the same constraint). SimpleQueue.put is + # lock-free/reentrant, so finalizers only enqueue; every public method drains the queue + # under the lock before doing its own work. + self._deferred_releases: queue.SimpleQueue[tuple[str, dict[str, torch.Tensor] | None]] = queue.SimpleQueue() self._entries: dict[str, _SharedWeightsEntry] = {} # Entries forgotten by `invalidate()` while still referenced by live cached models (e.g. a # locked, stale-marked cache entry mid-generation). They can no longer be acquired or peeked, @@ -76,6 +85,7 @@ def acquire(self, key: str, state_dict: dict[str, torch.Tensor]) -> dict[str, to re-pointing its module at these tensors and dropping the `state_dict` it passed in. """ with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) if entry is None: entry = _SharedWeightsEntry( @@ -95,6 +105,7 @@ def peek(self, key: str) -> dict[str, torch.Tensor] | None: itself increment the count. """ with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) return entry.state_dict if entry is not None else None @@ -102,6 +113,7 @@ def set_shell(self, key: str, shell: object) -> None: """Register the empty (meta-weight) structural clone for `key`, if an entry exists and none is set yet. A no-op when the key has no canonical entry (e.g. keep_ram_copy disabled).""" with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) if entry is not None and entry.shell is None: entry.shell = shell @@ -109,6 +121,7 @@ def set_shell(self, key: str, shell: object) -> None: def get_shell(self, key: str) -> object | None: """Return the registered meta-weight shell for `key`, or None if absent.""" with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) return entry.shell if entry is not None else None @@ -127,21 +140,47 @@ def release(self, key: str, state_dict: dict[str, torch.Tensor] | None = None) - lets go. """ with self._lock: - entry = self._entries.get(key) - if entry is not None and (state_dict is None or entry.state_dict is state_dict): - entry.refcount -= 1 - if entry.refcount <= 0: - del self._entries[key] + self._drain_deferred_locked() + self._release_locked(key, state_dict) + + def release_deferred(self, key: str, state_dict: dict[str, torch.Tensor] | None = None) -> None: + """Post a release to be applied by the next store operation, WITHOUT taking the store lock. + + This is the only release entry point that is safe from GC context (weakref.finalize + callbacks, __del__): it only enqueues. A finalizer can fire at any allocation point on any + thread — including while that same thread holds self._lock inside acquire() — so taking the + non-reentrant lock here could self-deadlock the process. Every public method drains the + queue under the lock, so the released bytes disappear from the accounting no later than the + next store operation (in particular, the next `total_bytes_in_use()` / budget query). + """ + self._deferred_releases.put((key, state_dict)) + + def _drain_deferred_locked(self) -> None: + """Apply all pending deferred releases. Caller must hold self._lock.""" + while True: + try: + key, state_dict = self._deferred_releases.get_nowait() + except queue.Empty: return - # Not the live canonical for `key` — it may be a retired (invalidated) entry whose - # tensors are still being counted against the RAM budget. - if state_dict is not None: - for i, retired in enumerate(self._retired): - if retired.state_dict is state_dict: - retired.refcount -= 1 - if retired.refcount <= 0: - del self._retired[i] - return + self._release_locked(key, state_dict) + + def _release_locked(self, key: str, state_dict: dict[str, torch.Tensor] | None) -> None: + """The body of release(). Caller must hold self._lock.""" + entry = self._entries.get(key) + if entry is not None and (state_dict is None or entry.state_dict is state_dict): + entry.refcount -= 1 + if entry.refcount <= 0: + del self._entries[key] + return + # Not the live canonical for `key` — it may be a retired (invalidated) entry whose + # tensors are still being counted against the RAM budget. + if state_dict is not None: + for i, retired in enumerate(self._retired): + if retired.state_dict is state_dict: + retired.refcount -= 1 + if retired.refcount <= 0: + del self._retired[i] + return def invalidate(self, model_key: str) -> int: """Forget the canonical entries (and shells) for `model_key` and all of its submodels, so no @@ -158,6 +197,7 @@ def invalidate(self, model_key: str) -> int: """ prefix = f"{model_key}:" with self._lock: + self._drain_deferred_locked() doomed = [key for key in self._entries if key == model_key or key.startswith(prefix)] for key in doomed: entry = self._entries.pop(key) @@ -169,11 +209,13 @@ def invalidate(self, model_key: str) -> int: def __contains__(self, key: str) -> bool: with self._lock: + self._drain_deferred_locked() return key in self._entries def refcount(self, key: str) -> int: """Return the current refcount for `key`, or 0 if not present.""" with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) return entry.refcount if entry is not None else 0 @@ -185,6 +227,7 @@ def total_bytes_in_use(self) -> int: it — i.e. the true RAM footprint of cached weights, not the per-device double-count. """ with self._lock: + self._drain_deferred_locked() return sum(entry.total_bytes for entry in self._entries.values()) + sum( entry.total_bytes for entry in self._retired ) @@ -192,10 +235,12 @@ def total_bytes_in_use(self) -> int: def retired_bytes(self) -> int: """Return the total size (in bytes) of retired (invalidated but still referenced) entries.""" with self._lock: + self._drain_deferred_locked() return sum(entry.total_bytes for entry in self._retired) def keys(self) -> list[str]: with self._lock: + self._drain_deferred_locked() return list(self._entries.keys()) diff --git a/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py b/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py index 79a34ff0d96..cfd9a8da501 100644 --- a/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py +++ b/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py @@ -113,9 +113,10 @@ def load_state_dict(self, *args, **kwargs): # type: ignore[override] def test_acquire_is_released_if_repoint_fails(): - # First device registers the canonical weights (refcount 1). + # First device registers the canonical weights (refcount 1). The wrapper must stay bound: an + # abandoned wrapper's collection-time finalizer releases its reference (by design). store = SharedCpuWeightsStore() - CachedModelWithPartialLoad(DummyModule(), CPU, keep_ram_copy=True, shared_store=store, cache_key="m") + first = CachedModelWithPartialLoad(DummyModule(), CPU, keep_ram_copy=True, shared_store=store, cache_key="m") assert store.refcount("m") == 1 # Second device adopts the canonical copy, but its re-point throws. The just-acquired reference @@ -124,3 +125,4 @@ def test_acquire_is_released_if_repoint_fails(): CachedModelWithPartialLoad(_RepointFailsModule(), CPU, keep_ram_copy=True, shared_store=store, cache_key="m") assert store.refcount("m") == 1 # back to just the first device, not leaked at 2 + assert first.uses_shared_weights # keep the first wrapper alive through the assertions above diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py index 656a10be534..b4f465edaa8 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py @@ -13,6 +13,7 @@ from unittest.mock import MagicMock, patch import pytest +import torch from invokeai.backend.model_manager.load.load_base import LoadedModelWithoutConfig from invokeai.backend.model_manager.load.model_cache import model_cache as model_cache_module @@ -91,6 +92,150 @@ def test_shared_model_counts_once_in_global_budget(mock_logger): cache_b.shutdown() +def _collect_until(predicate, attempts: int = 5) -> bool: + """gc.collect() until predicate() holds — finalizer chains can need more than one pass.""" + for _ in range(attempts): + gc.collect() + if predicate(): + return True + return predicate() + + +def test_shutdown_releases_shared_weights_synchronously(mock_logger): + """shutdown() must release its resident records' shared references itself: the finalizer + fallback only enqueues, and at teardown there may be no later store operation to drain the + queue — so relying on collection would leave the canonical tensors pinned indefinitely.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + assert store.refcount("m") == 1 + + cache.shutdown() + # No gc, no further store activity needed: the release was synchronous. + assert store._entries.get("m") is None + assert store._deferred_releases.qsize() == 0 + + +def test_dropped_cache_releases_shared_weights_on_collection(mock_logger): + """A cache dropped without shutdown() must not strand its shared-weights references: + the store's refcount and bytes — and therefore the budget total — must return to zero once the + cache is collected.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) # keep_ram_copy=True -> shared weights + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + cache_ref = weakref.ref(cache) + del cache + assert _collect_until(lambda: cache_ref() is None) + assert store.refcount("m") == 0 + assert store.total_bytes_in_use() == 0 + assert budget.total_in_use() == 0 + + +def test_collection_release_is_deferred_not_taken_under_the_store_lock(mock_logger): + """The collection-time release must only ENQUEUE: it runs in GC context, where taking the + store's non-reentrant lock (e.g. while another frame on the same thread is inside acquire()) + would self-deadlock the process. The queue is drained by the next store operation.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + + cache_ref = weakref.ref(cache) + del cache + assert _collect_until(lambda: cache_ref() is None) + # The finalizer has fired, but it must not have touched the entries directly: the refcount is + # still 1 when read without the public (draining) API, and the release sits in the queue. + assert store._entries["m"].refcount == 1 + assert store._deferred_releases.qsize() == 1 + # The next public operation applies it. + assert store.refcount("m") == 0 + assert store._deferred_releases.qsize() == 0 + + +def test_normal_eviction_and_collection_release_exactly_once(mock_logger): + """An entry evicted through _delete_cache_entry (which calls release_shared_weights) must not + be released AGAIN when its wrapper is later collected — the second device's reference would be + freed out from under it.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache_a = _make_cache(store, budget, mock_logger) + cache_b = _make_cache(store, budget, mock_logger) + cache_a.put("m", DummyModule()) + cache_b.put("m", DummyModule()) + _use_and_release(cache_a, "m") + _use_and_release(cache_b, "m") + assert store.refcount("m") == 2 + + # Normal eviction on cache_a releases its reference synchronously (and detaches the fallback). + assert cache_a.evict_unlocked_for_peer(lambda: False) == 1 + assert store.refcount("m") == 1 + + # Collecting cache_a afterwards must not decrement again on cache_b's behalf. + ref_a = weakref.ref(cache_a) + del cache_a + assert _collect_until(lambda: ref_a() is None) + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + ref_b = weakref.ref(cache_b) + del cache_b + assert _collect_until(lambda: ref_b() is None) + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_dropped_cache_releases_a_retired_shared_entry(mock_logger): + """invalidate() moves a still-referenced entry to the retired list, matched later by state-dict + identity. A holder that is collected (rather than evicted) must still free the retired entry's + accounting via the deferred release.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + assert store.invalidate("m") == 1 + assert store.retired_bytes() == S + + cache_ref = weakref.ref(cache) + del cache + assert _collect_until(lambda: cache_ref() is None) + assert store.retired_bytes() == 0 + assert store.total_bytes_in_use() == 0 + assert budget.total_in_use() == 0 + + +def test_collected_partial_load_wrapper_releases_shared_weights(mock_logger): + """CachedModelWithPartialLoad (the partial-loading wrapper) has the same collection-time + release as CachedModelOnlyFullLoad.""" + from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import ( + CachedModelWithPartialLoad, + ) + + store = SharedCpuWeightsStore() + wrapped = CachedModelWithPartialLoad( + model=DummyModule(), + compute_device=torch.device("cpu"), + keep_ram_copy=True, + shared_store=store, + cache_key="m", + ) + assert store.refcount("m") == 1 + + wrapped_ref = weakref.ref(wrapped) + del wrapped + assert _collect_until(lambda: wrapped_ref() is None) + assert store.refcount("m") == 0 + assert store.total_bytes_in_use() == 0 + + def test_non_shared_model_counts_per_device(mock_logger): store = SharedCpuWeightsStore() budget = RamBudget(max_bytes=10**12, shared_store=store)