""" Disk-based dataset cache for DecoupleRpy. WHY DISK INSTEAD OF MEMORY --------------------------- The MCP server runs as a short-lived subprocess: every tool call spawns a new `python server.py` process and exits after returning the result. A module-level memory dict is destroyed at the end of each call. /tmp on HuggingFace Spaces persists for the lifetime of the container (hours to days), so caching parsed h5ad files there eliminates the ~50MB GEO download on all calls after the first. Cache layout ------------ /tmp/decoupleRpy/cache/ GSE71729.h5ad ← parsed AnnData (fast to read_h5ad, no download) GSE71729.gpl_cache/ GPL14951.pkl ← probe→gene mapping dict """ from __future__ import annotations import json import re import threading from pathlib import Path from typing import Any CACHE_DIR = Path("/tmp/decoupleRpy/cache") # ── In-memory AnnData cache (process-lifetime) ─────────────────────────────── # When server.py runs as a persistent HTTP process (`--transport http`), the # same h5ad path is read on nearly every tool call of a multi-step run. Parsing # a 50MB h5ad from disk each time is pure waste once the process is long-lived. # This module-level cache holds the parsed AnnData for the process lifetime, # keyed by (abspath, mtime, size) so a rewritten file is transparently # re-read. Each read returns a `.copy()`, so callers may freely mutate (and # write back) without corrupting the shared, read-mostly cache — this is what # makes it safe to share one resident server across concurrent Gradio sessions. # Under stdio transport (one process per call) it simply never gets a second # hit, so it's a no-op there rather than a correctness risk. _MEM_LOCK = threading.Lock() _MEM_CACHE: dict[str, Any] = {} _MEM_MAX_ENTRIES = 8 # bound memory; bulk h5ads are large def _mem_key(path: str) -> str | None: try: st = Path(path).stat() except OSError: return None return f"{Path(path).resolve()}::{int(st.st_mtime)}::{st.st_size}" def read_h5ad_cached(path: str) -> Any: """Read an h5ad, reusing a process-lifetime in-memory parse when possible. Returns a fresh `.copy()` every call, so the caller owns the object and may mutate or write it back without affecting the cache. Falls back to a plain `sc.read_h5ad` if the path can't be stat'd (e.g. a URL or missing file) so the caller sees the normal error. """ import scanpy as sc key = _mem_key(path) if key is None: return sc.read_h5ad(path) with _MEM_LOCK: cached = _MEM_CACHE.get(key) if cached is not None: return cached.copy() adata = sc.read_h5ad(path) with _MEM_LOCK: # Drop stale entries for the same resolved path (older mtime/size). resolved = str(Path(path).resolve()) for k in [k for k in _MEM_CACHE if k.startswith(resolved + "::") and k != key]: _MEM_CACHE.pop(k, None) _MEM_CACHE[key] = adata # Evict oldest if over the bound (dict preserves insertion order). while len(_MEM_CACHE) > _MEM_MAX_ENTRIES: _MEM_CACHE.pop(next(iter(_MEM_CACHE))) return adata.copy() # ── Cache key helpers ──────────────────────────────────────────────────────── def url_to_cache_key(url_or_path: str) -> str: """Return a stable GSE accession key from any URL or path string. Picks the match with the most digits to avoid partial matches like GSE71 from GSE71nnn/GSE71729.""" matches = re.findall(r"GSE(\d+)", url_or_path, re.IGNORECASE) if not matches: return url_or_path return "GSE" + max(matches, key=len) # ── AnnData (dataset) cache ────────────────────────────────────────────────── def _dataset_cache_path(cache_key: str) -> Path: return CACHE_DIR / f"{cache_key}.h5ad" def get_dataset(cache_key: str) -> Any | None: """Return a cached AnnData, or None if not cached.""" path = _dataset_cache_path(cache_key) if path.exists() and path.stat().st_size > 0: try: import scanpy as sc print(f"[cache] Disk hit for {cache_key} — loading from {path}") return sc.read_h5ad(path) except Exception as exc: print(f"[cache] Disk read failed for {cache_key}: {exc} — will re-download") path.unlink(missing_ok=True) return None def set_dataset(cache_key: str, adata: Any) -> None: """Write an AnnData to the disk cache.""" try: CACHE_DIR.mkdir(parents=True, exist_ok=True) path = _dataset_cache_path(cache_key) adata.write_h5ad(path) print(f"[cache] Wrote {cache_key} to disk cache ({path.stat().st_size // 1024} KB)") except Exception as exc: print(f"[cache] Failed to write {cache_key} to disk cache: {exc}") def is_loaded(cache_key: str) -> bool: return _dataset_cache_path(cache_key).exists() # ── GPL probe mapping cache ────────────────────────────────────────────────── def _gpl_cache_path(accession: str) -> Path: # JSON (not pickle): the mapping is a plain dict[str, str], so JSON is safe # to deserialize and human-inspectable. Old .pkl files are simply ignored # (cache miss → regenerated). return CACHE_DIR / "gpl" / f"{accession}.json" def get_gpl(accession: str) -> dict[str, str] | None: """Return a cached probe→gene mapping, or None if not cached.""" path = _gpl_cache_path(accession) if path.exists() and path.stat().st_size > 0: try: print(f"[cache] Disk hit for GPL {accession}") with open(path, "r", encoding="utf-8") as f: return json.load(f) except Exception as exc: print(f"[cache] Disk read failed for GPL {accession}: {exc}") path.unlink(missing_ok=True) return None def set_gpl(accession: str, mapping: dict[str, str]) -> None: """Write a probe→gene mapping to the disk cache.""" try: _gpl_cache_path(accession).parent.mkdir(parents=True, exist_ok=True) with open(_gpl_cache_path(accession), "w", encoding="utf-8") as f: json.dump(mapping, f) print(f"[cache] Wrote GPL {accession} to disk cache") except Exception as exc: print(f"[cache] Failed to write GPL {accession} to disk cache: {exc}") # ── Preloader (optional warm-up, still useful for first-query speed) ───────── def preload_datasets() -> None: """ Pre-populate the disk cache at startup by downloading all registered geo_series_matrix datasets. Runs in a background thread so the server stays responsive. No-ops if the cache file already exists. """ import pandas as pd import scanpy as sc from src.datasets.registry import get_registry from src.workflows.geo import load_geo_series_matrix_lines, parse_geo_series_matrix_lines from src.tools.rna import decode_geo_numeric_codes registry = get_registry() for dataset_id in registry.list(): raw = registry.get(dataset_id) if raw is None: continue expr_src = raw.get("expression_source", {}) if expr_src.get("type") != "geo_series_matrix": continue url = expr_src.get("url") if not url: continue cache_key = url_to_cache_key(url) if is_loaded(cache_key): print(f"[cache] {dataset_id} already on disk — skipping preload") continue print(f"[cache] Preloading {dataset_id} from {url} …") try: lines = load_geo_series_matrix_lines(url) parsed = parse_geo_series_matrix_lines(lines) sample_ids = parsed["sample_ids"] probe_ids = parsed["probe_ids"] X = parsed["X"] sample_characteristics = parsed["sample_characteristics"] obs_df = pd.DataFrame(index=sample_ids) for key, values in sample_characteristics.items(): col_name = key.lower().replace(" ", "_").replace("-", "_") if len(values) == len(sample_ids): obs_df[col_name] = values for col in list(obs_df.columns): result = decode_geo_numeric_codes(col, obs_df[col]) if result is not None: base_name, decoded = result if base_name not in obs_df.columns: obs_df[base_name] = decoded adata = sc.AnnData(X=X, obs=obs_df) adata.var.index = pd.Index(probe_ids, name="probe_id") set_dataset(cache_key, adata) print(f"[cache] {dataset_id} preloaded: {adata.n_obs} × {adata.n_vars}") except Exception as exc: print(f"[cache] Failed to preload {dataset_id}: {exc}")