Skip to content

Space's unbounded method caches retain every compared Space for the process lifetime #4

Description

@maxnus

Summary

Space._distance_from_zero_singular_value and Space._distance_from_one_eigenvalue in src/btensor/space.py are decorated with functools.cache, which is unbounded. Because the decorator is applied to a method, each cache entry holds a strong reference to both self and the other argument. Every Space that takes part in at least one comparison is therefore retained for the life of the process, along with the Basis it wraps and that basis's array data.

@cache  # noqa: B019
def _distance_from_zero_singular_value(self, other: Space) -> float:
    ...

@cache  # noqa: B019
def _distance_from_one_eigenvalue(self, other: Space) -> float:
    ...

These are reached from __eq__, __ne__, __lt__, __le__, __gt__, __ge__ and __or__, so an ordinary comparison is enough to pin a Space.

Noticed while clearing ruff findings in #3. The noqa: B019 comments there acknowledge the pattern but do not address the retention; this issue is the follow-up for the actual decision.

Why it matters

The growth is proportional to the number of distinct pairs compared, not to the number of live objects, so a long-running process that builds many transient bases — a loop over geometries, a parameter scan, a notebook session — accumulates memory that never comes back. The arrays hanging off the retained bases are the expensive part, not the float results.

This is not a correctness bug and there is no reported symptom. Space.__hash__ returns self._basis.id, so the cache keys are stable and the cached values are correct.

Prior art in this repo

Basis.get_transformation has the same shape but is already bounded, with the tradeoff written down:

_transformation_cache_size = 128

# B019: lru_cache on a method keeps `self` alive for as long as the entry is
# ...
@lru_cache(_transformation_cache_size)  # noqa: B019

Whatever is decided here, matching that precedent would at least make the two consistent.

Options

  1. Bound them — swap @cache for @lru_cache(maxsize=...), mirroring Basis. Smallest change, caps the growth, but still retains up to maxsize spaces and the eviction policy is arbitrary with respect to what is actually live.
  2. Move the cache onto the instance — a per-Space dict keyed by other's basis id, so entries die with the Space that owns them. Removes the self leak; other is still pinned unless the values are keyed by id and held weakly.
  3. Key on basis ids and hold the spaces weakly — a WeakValueDictionary or an explicit id-keyed cache. Most correct, most work, and needs care because basis ids must not be recycled.
  4. Drop the caching — measure first. If the SVD and eigendecomposition are not actually hot in realistic use, this is the simplest resolution.

Option 1 is the cheap fix and option 4 is the honest one; picking between them really wants a measurement of how often these are recomputed in practice.

Deliberately not decided here

Changing this is a performance tradeoff, not a cleanup, so it was left alone in #3 rather than guessed at. Whoever picks it up should know the cache is doing real work — both methods run a decomposition — so removing or shrinking it has a cost that should be measured rather than assumed.

Aside

While in this file: the try: from functools import cache / except ImportError: from functools import lru_cache as cache fallback at the top is now dead code, since requires-python is >=3.11 as of #3. Worth deleting in the same change.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions