feat: process-wide model weight cache (v0.1.5)

[AzulGarza]

Summary

• Add a process-wide LRU model weight cache (default max_cached_models=1) so local weight-loading forecasters reuse loaded checkpoints across repeated forecast() calls instead of reloading on every request.
• Wire all weight-loading models through the base Forecaster._cached_model() helper; every forecaster accepts reuse_loaded_model and exposes clear_model_cache().
• Export set_max_cached_models() from the package root; FoundationForecast(clean_cache=True) clears cached weights after each model in multi-model runs.
• Document usage in a new Model Weight Cache guide and ship v0.1.5 release notes.

Test plan

• uv run pytest tests/core/test_model_weight_cache.py tests/core/ tests/test_foundation_forecast.py
• uv run ruff check foundationforecast
• Merge PR
• Tag v0.1.5 on main to trigger PyPI release workflow

Made with Cursor

[Copilot]

🟡 Changes recommended

Unresolved cache concurrency and cleanup defects, plus stale downstream lock metadata, block approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds a process-wide LRU cache for locally loaded model weights, with reuse controls, cleanup APIs, tests, documentation, and v0.1.5 release metadata.

Changes:

• Adds shared model caching and configuration.
• Integrates caching across supported forecasters.
• Updates tests, documentation, changelogs, and package versioning.

File summaries

File	Summary
uv.lock	Updates lock metadata; critical (1 vote): downstream GIFT-Eval lock metadata remains stale.
tests/helpers.py	Updates test forecasters for cache options.
tests/core/test_model_weight_cache.py	Adds cache behavior and cleanup tests.
pyproject.toml	Bumps version to 0.1.5; critical (2 votes): GIFT-Eval frozen installation may reject stale metadata.
mkdocs.yml	Adds cache documentation navigation.
foundationforecast/models/toto.py	Integrates Toto caching; moderate (1 vote): cleanup can trigger network I/O for unforecast remote models.
foundationforecast/models/tirex.py	Integrates TiRex model caching.
foundationforecast/models/timesfm.py	Integrates TimesFM caching; moderate (1 vote): prediction_length causes unnecessary cache misses.
foundationforecast/models/timegpt.py	Integrates TimeGPT caching.
foundationforecast/models/tafsut.py	Integrates TAFSUT caching.
foundationforecast/models/tabpfn.py	Integrates TabPFN caching.
foundationforecast/models/t0.py	Integrates T0 caching.
foundationforecast/models/sundial.py	Integrates Sundial caching.
foundationforecast/models/patchtst_fm.py	Integrates PatchTST-FM caching.
foundationforecast/models/moirai.py	Integrates Moirai caching.
foundationforecast/models/flowstate.py	Integrates FlowState caching.
foundationforecast/models/chronos.py	Integrates Chronos caching.
foundationforecast/core/model_weight_cache.py	Implements the LRU cache; critical (2 votes): cache operations are unsynchronized; moderate (1 vote): zero capacity bypasses cleanup.
foundationforecast/core/gluonts_forecaster.py	Adds cached predictor support.
foundationforecast/core/forecaster.py	Adds shared cache APIs.
foundationforecast/core/cached_forecaster.py	Provides cache context helpers; moderate (1 vote): zero-capacity loads bypass release_model.
foundationforecast/_foundation_forecast.py	Adds cleanup integration; moderate (2 votes): fallback_model is excluded from cleanup.
foundationforecast/__init__.py	Exports cache configuration.
docs/model-weight-cache.md	Documents cache usage and configuration.
docs/model-hub.md	Links to cache documentation.
docs/changelogs/v0.1.5.md	Adds v0.1.5 release notes.
docs/changelogs/index.md	Registers the new changelog.
docs/api/models/utils/forecaster.md	Documents forecaster cache APIs.

Review details

Suppressed comments (7)

foundationforecast/core/cached_forecaster.py:28

• When max_cached_models is 0, this returns a freshly loaded model directly, bypassing the finally in cached_model_context; set_max_cached_models(0) therefore never calls release_model for forecast-created models. That leaves the CUDA/MPS allocator cleanup path unused, contrary to the documented “always load + release” behavior. Handle zero-capacity as an uncached load with the same try/finally release path.

        yield get_model_weight_cache().get_or_load(cache_key, loader)

foundationforecast/core/model_weight_cache.py:38

• When the cache is full, the new model is loaded before the least-recently-used model is evicted. With the default capacity of 1, switching models therefore keeps both checkpoints resident during loader(), so loading a second large GPU model can OOM before the eviction runs. Evict and release the old entry before loading a replacement (while preserving the old entry if the load fails, if that behavior is required).

        model = loader()
        self._cache[key] = model
        while len(self._cache) > self._max_cached_models:
            _, evicted = self._cache.popitem(last=False)
            release_model(evicted)

foundationforecast/core/model_weight_cache.py:64

• del model only deletes the local parameter in release_model; it cannot drop the caller's model/evicted reference. The subsequent empty_cache() calls therefore run while the model is still held (and clear() still retains entries until after its loop), so eviction and clear_model_cache() may not actually release GPU allocations as documented. Drop the owning references before invoking the device-cache cleanup.

def release_model(model: Any) -> None:
    del model
    try:
        import torch

foundationforecast/core/model_weight_cache.py:54

• clear_prefix() uses an unbounded startswith, so clearing a forecaster whose prefix is T0:repo also clears cached keys such as T0:repo2. This can make clear_model_cache() evict an unrelated forecaster's weights; match the prefix exactly or require the separator after it.

        keys = [key for key in self._cache if key.startswith(prefix)]
        for key in keys:
            self.clear(key)

foundationforecast/core/model_weight_cache.py:30

• When max_cached_models == 0, this returns a freshly loaded model without placing it in the cache. However, cached_model_context() still takes its cached branch for reuse_loaded_model=True, so it never enters the finally: release_model(...) path; the documented zero-cache mode therefore bypasses model cleanup. Treat zero capacity as the non-cached context path and add a regression test.

    def get_or_load(self, key: str, loader: Callable[[], T]) -> T:
        if self._max_cached_models == 0:
            return loader()

foundationforecast/models/timesfm.py:310

• _load_predictor() and its ModelConfig do not use prediction_length, so including it in this cache key creates a separate evaluator for every horizon. With the default one-entry LRU, alternating h values reloads the same TimesFM 3 checkpoint each time; omit this argument from the v3 key (or cache only the underlying weights).

    def _model_cache_key(self, prediction_length: int) -> str:
        kwargs_key = tuple(sorted((self.kwargs or {}).items()))
        return (
            f"{self._model_cache_prefix()}:"
            f"{self.context_length}:{self.batch_size}:"
            f"{prediction_length}:{kwargs_key}"

foundationforecast/models/toto.py:152

• clear_model_cache() calls _model_cache_prefix(), so FoundationForecast(clean_cache=True) invokes _is_toto2() for Toto instances that have not forecast yet. For a remote repo this performs hf_hub_download(config.json), making cleanup perform network I/O and potentially fail before that model is run. Keep prefix resolution side-effect-free or clear only keys that have already been created.

    def _model_cache_prefix(self) -> str | None:
        version = "v2" if self._is_toto2() else "v1"
        return f"{type(self).__qualname__}:{self.repo_id}:{version}"

• Files reviewed: 27/28 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Unresolved critical issues can permit duplicate concurrent checkpoint loads and retain evicted weights during device-cache cleanup, risking out-of-memory failures.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

foundationforecast/core/model_weight_cache.py:97

• The singleton is initialized with an unsynchronized check-and-create. If the first forecasts arrive concurrently, two ModelWeightCache instances can be constructed and the later assignment can discard the first, so requests may load duplicate checkpoints and the process-wide cache loses one. Protect initialization with a module-level lock or initialize the cache eagerly.

    if _model_weight_cache is None:
        _model_weight_cache = ModelWeightCache()
    return _model_weight_cache

• Files reviewed: 29/31 changed files
• Comments generated: 2
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Unresolved concurrency, eviction, and failure-path cleanup issues may cause duplicate loads or retained GPU memory.

Review details

Suppressed comments (5)

foundationforecast/core/model_weight_cache.py:112

• This clear path only removes entries already in _cache; an in-flight load is neither cancelled nor invalidated. If clear_model_cache() runs while another thread is loading the same key, that loader inserts the model after clear() returns, leaving weights cached despite the clear request. Coordinate clearing with _inflight loads or use a generation/invalidation check.

    def clear(self, key: str | None = None) -> None:
        with self._lock:
            if key is None:
                models = list(self._cache.values())
                self._cache.clear()

foundationforecast/core/model_weight_cache.py:148

• release_model() only deletes its helper parameter, while callers still retain the object in the models/evicted lists (and the non-cached context manager still has its local). _empty_device_cache() therefore runs before the tensors become unreachable, so an LRU eviction or clear_model_cache() can return without actually flushing freed GPU blocks. Drop the owning references before performing device cleanup, or move cleanup out of this helper.

def release_model(model: Any) -> None:
    del model
    _empty_device_cache()

foundationforecast/core/model_weight_cache.py:106

• A waiter only rechecks _cache after the event is signaled. With max_cached_models=1, another key can finish between the original load's insertion and this lock acquisition, evicting this key; the waiter then loops and invokes loader() a second time, so concurrent requests for the same key can load duplicate models under contention. Return inflight.model when the load completed before retrying the cache lookup.

                cached = self._cache.get(key)
                if cached is not None:
                    self._cache.move_to_end(key)
                    return cached

foundationforecast/core/model_weight_cache.py:68

• The LRU entry is evicted and released before loader() runs. If the new checkpoint download or initialization raises, the previously warm model is lost even though no replacement was loaded, so the next request must reload it. Defer eviction until the successful insertion path below (or restore the evicted entries when loading fails) to preserve a usable cache on transient load errors.

                    evicted = self._evict_if_needed_locked()
                    inflight = _InflightLoad()
                    self._inflight[key] = inflight
                    is_loader = True
                else:
                    is_loader = False

            for model in evicted:
                release_model(model)

foundationforecast/core/multi_model.py:30

• With the shared cache, a model that raises before res_df_model is produced keeps its successfully loaded weights cached, but _call_models only invokes _clean_model_cache() after the try/merge block. Thus clean_cache=True does not release weights when a primary model fails without a fallback (or when the fallback also fails), so an unsuccessful multi-model run can retain GPU memory. Put the cleanup in a finally for each model attempt while preserving fallback handling.

        models_to_clear = list(self.models)
        if self.fallback_model is not None:
            models_to_clear.append(self.fallback_model)
        for model in models_to_clear:

• Files reviewed: 29/31 changed files
• Comments generated: 0 new
• Review effort level: Lite

[codspeed]

Merging this PR will improve performance by ×2.6

⚡ 3 improved benchmarks
✅ 5 untouched benchmarks

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	WallTime	test_model_forecast[chronos-bolt]	675.1 ms	119.3 ms	×5.7
⚡	WallTime	test_model_forecast[timesfm-1]	6.1 s	2.7 s	×2.3
⚡	WallTime	test_model_forecast[moirai-1.1]	1.9 s	1.4 s	+41.19%

Tip

Curious why performance improved? Comment @codspeedbot explain why performance improved on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing feat/model-weight-cache (af6b01a) with main (8e5e745)}

[AzulGarza]

@codspeedbot explain why performance improved

[codspeed]

@codspeedbot explain why performance improved

Why performance improved

Short answer: the benchmarks are no longer paying a checkpoint load on every forecast() call. The gain you see is almost exactly the per-call weight-loading cost that this PR removes.

What changed in the measured path

On main, every model's _get_model() / _get_predictor() context manager loaded the checkpoint on entry and threw it away on exit:

model = cls.from_pretrained(self.repo_id, device_map=..., torch_dtype=self.dtype)
try:
    yield model
finally:
    del model
    torch.cuda.empty_cache()

Since forecast() opens that context on each call, every call paid checkpoint deserialization + module construction (plus the Hugging Face metadata lookups such as repo_exists / config reads in the TimesFM and Toto paths). On this branch, the same loaders are routed through Forecaster._cached_model() → ModelWeightCache.get_or_load() (process-wide LRU, max_cached_models=1, keyed by model config), so the load happens once per process instead of once per call.

Why that shows up so strongly in the benchmarks

tests/benchmarks/test_models.py benchmarks model.forecast with session-scoped model fixtures, and pytest-codspeed runs warmup iterations followed by several measured rounds in the same process. On this branch the cache is already warm after warmup, so every measured round is a cache hit; on main every round re-loaded the weights. The reported delta is therefore a direct estimate of the removed per-call load cost:

Benchmark	BASE (min)	HEAD (min)	Per-call load removed
test_model_forecast[chronos-bolt]	675.1 ms	119.3 ms	~556 ms
test_model_forecast[timesfm-1]	6.14 s	2.66 s	~3.5 s
test_model_forecast[moirai-1.1]	1.94 s	1.37 s	~0.57 s

The ranking matches load weight vs. inference weight: chronos-bolt-tiny inference is tiny (~119 ms), so loading was ~85% of the call → ×5.7. timesfm-1.0-200m pays a large checkpoint construction → ×2.3. moirai-1.1-R-small spends most of its time in GluonTS inference, so removing MoiraiModule.from_pretrained + create_predictor only buys +41%.

test_model_forecast[toto] is correctly reported as untouched: it is cached too, but the call is ~38 s of inference (41.1 s → 38.3 s) and only one round is measured, so the load saving is inside the noise band rather than a detectable change.

Two things worth noting

• Not an environment artifact. The comparison did flag a CPU change (EPYC 7763 → 9V74), but that only affected the four Simulation benchmarks in tests/benchmarks/test_core.py; the WallTime model benchmarks ran on matching codspeed-macro runners, so the improvements are attributable to the code change.
• This measures the warm path. A single cold forecast() in a fresh process still pays the load once — the win is for repeated calls in a long-lived process, which is what the benchmark exercises. A useful side effect is that these benchmarks now mostly measure inference instead of I/O, making future regressions in the model code much easier to see. With max_cached_models=1, a benchmark that alternated between two models within one call would thrash the cache; the current suite is safe because each parametrization runs its rounds consecutively.

No code changes were needed for this request, so the branch is untouched. Flamegraphs were not available for these walltime benchmarks (callgraph generation reported a missing root frame), so the above is based on the walltime distributions and the diff of the measured code path.

Summary

• Warming up the performance wizard...
• Inspect CodSpeed runs for base and head — Compared head af6b01a vs base main 8e5e745; 3 improvements, 5 unchanged; walltime model benchmarks ran on matching runners.
• Analyze the PR changes and benchmark code — Confirmed base loaded checkpoints inside _get_model/_get_predictor on every forecast() call; head routes loads through the process-wide LRU cache.
• Compare profiles for the improved benchmarks — No flamegraphs available for these walltime benchmarks (callgraph generation reported MissingRootFrame), so the analysis is based on walltime distributions plus the code diff.
• Report the explanation — Explanation posted: the measured gain is the per-call checkpoint load removed from every forecast() call.
• Handling comment — Comment handling finished successfully

_{Was this response helpful? React with 👍 or 👎 to help us improve CodSpeedBot.}