feat(ui/server): add configurable LRU cache sizes for models, datasets, saes, and circuits

Author: dest1n1s

server/config.py

@@ -9,3 +9,10 @@
 client = MongoClient(MongoDBConfig())
 sae_series = os.environ.get("SAE_SERIES", "default")
 tokenizer_only = os.environ.get("TOKENIZER_ONLY", "false").lower() == "true"
+
+# LRU cache sizes (configurable via environment variables)
+LRU_CACHE_SIZE_SAMPLES = int(os.environ.get("LRU_CACHE_SIZE_SAMPLES", "128"))
+LRU_CACHE_SIZE_MODELS = int(os.environ.get("LRU_CACHE_SIZE_MODELS", "8"))
+LRU_CACHE_SIZE_DATASETS = int(os.environ.get("LRU_CACHE_SIZE_DATASETS", "16"))
+LRU_CACHE_SIZE_SAES = int(os.environ.get("LRU_CACHE_SIZE_SAES", "8"))
+LRU_CACHE_SIZE_CIRCUITS = int(os.environ.get("LRU_CACHE_SIZE_CIRCUITS", "64"))

server/logic/loaders.py

@@ -5,12 +5,20 @@
 from lm_saes.abstract_sae import AbstractSparseAutoEncoder
 from lm_saes.backend import LanguageModel
 from lm_saes.resource_loaders import load_dataset_shard, load_model
-from server.config import client, device, sae_series, tokenizer_only
+from server.config import (
+    LRU_CACHE_SIZE_DATASETS,
+    LRU_CACHE_SIZE_MODELS,
+    LRU_CACHE_SIZE_SAES,
+    client,
+    device,
+    sae_series,
+    tokenizer_only,
+)
 from server.utils.common import synchronized
 
 
 @synchronized
-@lru_cache(maxsize=8)
+@lru_cache(maxsize=LRU_CACHE_SIZE_MODELS)
 def get_model(*, name: str) -> LanguageModel:
     """Load and cache a language model."""
     cfg = client.get_model_cfg(name)
@@ -22,7 +30,7 @@ def get_model(*, name: str) -> LanguageModel:
 
 
 @synchronized
-@lru_cache(maxsize=16)
+@lru_cache(maxsize=LRU_CACHE_SIZE_DATASETS)
 def get_dataset(*, name: str, shard_idx: int = 0, n_shards: int = 1) -> Dataset:
     """Load and cache a dataset shard."""
     cfg = client.get_dataset_cfg(name)
@@ -31,7 +39,7 @@ def get_dataset(*, name: str, shard_idx: int = 0, n_shards: int = 1) -> Dataset:
 
 
 @synchronized
-@lru_cache(maxsize=8)
+@lru_cache(maxsize=LRU_CACHE_SIZE_SAES)
 def get_sae(*, name: str) -> AbstractSparseAutoEncoder:
     """Load and cache a sparse autoencoder."""
     path = client.get_sae_path(name, sae_series)

server/logic/samples.py

@@ -4,11 +4,11 @@
 import numpy as np
 
 from lm_saes.database import FeatureAnalysisSampling
-from server.config import client, sae_series
+from server.config import LRU_CACHE_SIZE_SAMPLES, client, sae_series
 from server.logic.loaders import get_dataset, get_model
 
 
-@lru_cache(maxsize=128)
+@lru_cache(maxsize=LRU_CACHE_SIZE_SAMPLES)
 def cached_extract_samples(
     sae_name: str,
     sae_series: str,
@@ -211,7 +211,13 @@ def list_feature_data(
             )
             samples = (
                 cached_extract_samples(
-                    sae_name, sae_series, feature.index, sampling.name, 0, sampling_size, visible_range=sampling_visible_range
+                    sae_name,
+                    sae_series,
+                    feature.index,
+                    sampling.name,
+                    0,
+                    sampling_size,
+                    visible_range=sampling_visible_range,
                 )
                 if sampling is not None
                 else []

server/routers/circuits.py

@@ -19,7 +19,7 @@
 from lm_saes.database import CircuitConfig, CircuitInput, CircuitStatus
 from lm_saes.lorsa import LowRankSparseAttention
 from lm_saes.sae import SparseAutoEncoder
-from server.config import client, sae_series
+from server.config import LRU_CACHE_SIZE_CIRCUITS, client, sae_series
 from server.logic.loaders import get_model, get_sae
 from server.logic.samples import list_feature_data
 from server.utils.common import make_serializable, synchronized
@@ -146,7 +146,7 @@ def concretize_graph_data(graph_data: dict[str, Any]):
 
 
 @synchronized
-@functools.lru_cache(maxsize=64)
+@functools.lru_cache(maxsize=LRU_CACHE_SIZE_CIRCUITS)
 def load_circuit_graph(*, circuit_id: str, node_threshold: float, edge_threshold: float) -> dict[str, Any]:
     """Load, prune, and concretize a circuit graph. Cached for repeated access."""
     circuit = client.get_circuit(circuit_id)