lorax-arg 0.1__py3-none-any.whl

lorax/cache/file_cache.py

@@ -0,0 +1,142 @@
+"""
+Unified File Cache for Lorax.
+
+Provides cached loading of FileContext objects which combine tree sequences,
+config, and metadata into a single eviction unit. When a file is evicted,
+all related data is evicted together.
+
+This replaces the separate tree_sequence, config, and metadata caches.
+"""
+
+import asyncio
+from pathlib import Path
+from typing import Optional
+
+import pandas as pd
+import tskit
+import tszip
+
+from lorax.cache.lru import LRUCacheWithMeta
+from lorax.cache.file_context import FileContext
+from lorax.constants import TS_CACHE_SIZE
+
+# Global lock for thread-safe cache operations
+_cache_lock = asyncio.Lock()
+
+# Global cache for FileContext objects
+# Uses LRUCacheWithMeta to track file mtime for cache validation
+_file_cache = LRUCacheWithMeta(max_size=TS_CACHE_SIZE)
+
+
+def _get_file_mtime(file_path: str) -> float:
+    """Get file modification time, or 0 if file doesn't exist."""
+    try:
+        return Path(file_path).stat().st_mtime
+    except (OSError, FileNotFoundError):
+        return 0.0
+
+
+def _load_tree_sequence(file_path: str):
+    """Load tree sequence from file based on extension."""
+    if file_path.endswith('.tsz'):
+        return tszip.load(file_path)
+    elif file_path.endswith('.trees'):
+        return tskit.load(file_path)
+    elif file_path.endswith('.csv'):
+        return pd.read_csv(file_path)
+    else:
+        raise ValueError(f"Unsupported file type: {file_path}")
+
+
+async def get_file_context(file_path: str, root_dir: str = None) -> Optional[FileContext]:
+    """
+    Get or load a FileContext for the given file path.
+
+    Validates mtime and returns cached context if valid.
+    Loads fresh if cache miss or file changed.
+
+    Args:
+        file_path: Path to the tree sequence file
+        root_dir: Root directory for relative paths (defaults to file's parent)
+
+    Returns:
+        FileContext object with tree_sequence, config, and metadata cache,
+        or None if file doesn't exist or fails to load.
+    """
+    # Import here to avoid circular dependency
+    from lorax.loaders.loader import compute_config
+
+    file_path_obj = Path(file_path)
+    if not file_path_obj.exists():
+        _file_cache.remove(file_path)
+        print(f"❌ File not found: {file_path}")
+        return None
+
+    current_mtime = _get_file_mtime(file_path)
+
+    # Double-checked locking optimization
+    # 1. Optimistic check (lock-free read)
+    ctx, cached_mtime = _file_cache.get_with_meta(file_path)
+    if ctx is not None:
+        if cached_mtime == current_mtime:
+            print(f"✅ Using cached FileContext: {file_path}")
+            return ctx
+        else:
+            print(f"🔄 File changed, reloading: {file_path}")
+            _file_cache.remove(file_path)
+
+    # Need to load - acquire lock
+    async with _cache_lock:
+        # 2. Check again under lock (in case another task loaded it while we waited)
+        ctx, cached_mtime = _file_cache.get_with_meta(file_path)
+        if ctx is not None and cached_mtime == current_mtime:
+            print(f"✅ Using cached FileContext (after lock): {file_path}")
+            return ctx
+
+        print(f"📂 Loading FileContext: {file_path}")
+
+        try:
+            # Load tree sequence
+            ts = await asyncio.to_thread(_load_tree_sequence, file_path)
+
+            # Compute config immediately (it's derived from ts)
+            effective_root_dir = root_dir or str(file_path_obj.parent)
+            config = await asyncio.to_thread(
+                compute_config, ts, file_path, effective_root_dir
+            )
+
+            # Create FileContext with empty metadata cache
+            ctx = FileContext(
+                file_path=file_path,
+                tree_sequence=ts,
+                config=config,
+                mtime=current_mtime
+            )
+
+            _file_cache.set(file_path, ctx, meta=current_mtime)
+            return ctx
+
+        except Exception as e:
+            print(f"❌ Failed to load {file_path}: {e}")
+            return None
+
+
+def get_file_cache_size() -> int:
+    """Return current number of cached files."""
+    return len(_file_cache.cache)
+
+
+# Backwards compatibility: provide get_or_load_ts that returns just the tree_sequence
+async def get_or_load_ts(file_path: str, root_dir: str = None):
+    """
+    Backwards-compatible function that returns just the tree sequence.
+
+    Prefer using get_file_context() directly for new code.
+    """
+    ctx = await get_file_context(file_path, root_dir)
+    return ctx.tree_sequence if ctx else None
+
+
+def get_ts_cache_size() -> int:
+    """Backwards-compatible alias for get_file_cache_size()."""
+    return get_file_cache_size()

lorax/cache/file_context.py

@@ -0,0 +1,72 @@
+"""
+FileContext: Unified cache entry for loaded files.
+
+Combines tree sequence, config, and metadata into a single eviction unit.
+When a file is evicted from cache, all related data is evicted together,
+preventing orphan metadata and ensuring atomic invalidation.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Optional, Union
+
+import pandas as pd
+import tskit
+
+from lorax.cache.lru import LRUCache
+
+
+@dataclass
+class FileContext:
+    """
+    Unified cache entry for a loaded file.
+
+    Combines tree sequence, config, and metadata into single
+    eviction unit. When file is evicted from cache, all related
+    data is evicted together.
+
+    Attributes:
+        file_path: Absolute path to the loaded file
+        tree_sequence: The loaded tskit.TreeSequence or pandas.DataFrame (for CSV)
+        config: Configuration dict computed from the tree sequence
+        mtime: File modification time at load (used for cache validation)
+
+    The nested _metadata cache stores per-key metadata lookups:
+        - "population" -> {sample_name: value}
+        - "region" -> {sample_name: value}
+        - "population:array" -> {arrow_buffer, unique_values, sample_node_ids}
+    """
+    file_path: str
+    tree_sequence: Union[tskit.TreeSequence, pd.DataFrame]
+    config: dict
+    mtime: float
+
+    # Per-key metadata cache (nested within file context)
+    # Keys: "population", "region", "population:array", etc.
+    _metadata: LRUCache = field(default_factory=lambda: LRUCache(max_size=10))
+
+    def get_metadata(self, key: str) -> Optional[Any]:
+        """Get cached metadata for a specific key."""
+        return self._metadata.get(key)
+
+    def set_metadata(self, key: str, value: Any) -> None:
+        """Cache metadata for a specific key."""
+        self._metadata.set(key, value)
+
+    def clear_metadata(self) -> None:
+        """Clear all cached metadata (e.g., when file reloaded)."""
+        self._metadata.clear()
+
+    @property
+    def is_csv(self) -> bool:
+        """Check if this is a CSV file (pandas DataFrame)."""
+        return isinstance(self.tree_sequence, pd.DataFrame)
+
+    @property
+    def is_tree_sequence(self) -> bool:
+        """Check if this is a tskit tree sequence."""
+        return isinstance(self.tree_sequence, tskit.TreeSequence)
+
+    @property
+    def ts(self) -> Union[tskit.TreeSequence, pd.DataFrame]:
+        """Alias for tree_sequence for backwards compatibility."""
+        return self.tree_sequence

lorax/cache/lru.py

@@ -0,0 +1,90 @@
+"""
+LRU Cache implementations for Lorax.
+
+Provides in-memory caching with eviction for:
+- Tree sequences (with mtime validation)
+- Configuration data
+- Metadata lookups
+"""
+
+from collections import OrderedDict
+
+
+class LRUCache:
+    """Simple LRU cache with eviction for large in-memory tskit/tszip objects."""
+    def __init__(self, max_size=5):
+        self.max_size = max_size
+        self.cache = OrderedDict()
+
+    def get(self, key):
+        if key in self.cache:
+            # Move to the end to mark as recently used
+            self.cache.move_to_end(key)
+            return self.cache[key]
+        return None
+
+    def set(self, key, value):
+        if key in self.cache:
+            # Update existing and mark as recently used
+            self.cache.move_to_end(key)
+        self.cache[key] = value
+        # Evict if too big
+        if len(self.cache) > self.max_size:
+            old_key, old_val = self.cache.popitem(last=False)
+            print(f"🧹 Evicted {old_key} from LRU cache to free memory")
+
+    def remove(self, key):
+        """Remove a specific key from the cache."""
+        if key in self.cache:
+            del self.cache[key]
+
+    def clear(self):
+        self.cache.clear()
+
+
+class LRUCacheWithMeta:
+    """
+    LRU cache with metadata support for cache validation.
+
+    Stores (value, metadata) tuples, allowing validation against
+    external state (e.g., file mtime) before returning cached values.
+    """
+    def __init__(self, max_size=5):
+        self.max_size = max_size
+        self.cache = OrderedDict()  # key -> (value, meta)
+
+    def get(self, key):
+        """Get value only (ignores metadata)."""
+        if key in self.cache:
+            self.cache.move_to_end(key)
+            return self.cache[key][0]
+        return None
+
+    def get_with_meta(self, key):
+        """Get (value, metadata) tuple."""
+        if key in self.cache:
+            self.cache.move_to_end(key)
+            value, meta = self.cache[key]
+            return value, meta
+        return None, None
+
+    def set(self, key, value, meta=None):
+        """Set value with optional metadata."""
+        if key in self.cache:
+            self.cache.move_to_end(key)
+        self.cache[key] = (value, meta)
+        # Evict if too big
+        if len(self.cache) > self.max_size:
+            old_key, (old_val, old_meta) = self.cache.popitem(last=False)
+            print(f"🧹 Evicted {old_key} from LRU cache to free memory")
+
+    def remove(self, key):
+        """Remove a specific key from the cache."""
+        if key in self.cache:
+            del self.cache[key]
+
+    def clear(self):
+        self.cache.clear()
+
+    def __len__(self):
+        return len(self.cache)

lorax/cache/tree_graph.py

@@ -0,0 +1,293 @@
+"""
+TreeGraph Cache for per-session caching of TreeGraph objects.
+
+Supports Redis for production (distributed) and in-memory for local mode.
+Enables efficient lineage and search operations by reusing constructed trees.
+"""
+
+import pickle
+import asyncio
+from collections import OrderedDict
+from typing import Dict, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from lorax.tree_graph import TreeGraph
+
+# TTL for Redis entries (matches session lifetime)
+REDIS_TTL_SECONDS = 7 * 24 * 60 * 60  # 7 days
+
+
+class TreeGraphCache:
+    """
+    Per-session cache for TreeGraph objects with Redis/local mode switching.
+
+    In production mode (with Redis), TreeGraph objects are serialized via pickle
+    and stored with per-session keys. In local mode, uses in-memory dict.
+
+    Features:
+    - Visibility-based eviction: trees not in visible set are evicted
+    - Thread-safe async operations
+    - Automatic mode detection based on Redis availability
+    """
+
+    def __init__(self, redis_client=None):
+        """
+        Initialize the cache.
+
+        Args:
+            redis_client: Optional async Redis client. If None, uses in-memory mode.
+        """
+        self.redis = redis_client
+        # Local cache: session_id -> OrderedDict{tree_index -> TreeGraph}
+        # OrderedDict maintains insertion order for LRU eviction
+        self._local_cache: Dict[str, OrderedDict] = {}
+        self._lock = asyncio.Lock()
+
+        mode = "Redis" if redis_client else "in-memory"
+        print(f"TreeGraphCache initialized in {mode} mode")
+
+    def _redis_key(self, session_id: str, tree_index: int) -> str:
+        """Generate Redis key for a cached TreeGraph."""
+        return f"treegraph:{{{session_id}}}:{tree_index}"
+
+    def _redis_session_pattern(self, session_id: str) -> str:
+        """Generate Redis pattern to match all trees for a session."""
+        return f"treegraph:{{{session_id}}}:*"
+
+    async def get(self, session_id: str, tree_index: int) -> Optional["TreeGraph"]:
+        """
+        Retrieve a cached TreeGraph.
+
+        Args:
+            session_id: Session identifier
+            tree_index: Tree index in the tree sequence
+
+        Returns:
+            TreeGraph if cached, None otherwise
+        """
+        if self.redis:
+            return await self._redis_get(session_id, tree_index)
+        else:
+            return self._local_get(session_id, tree_index)
+
+    async def set(
+        self,
+        session_id: str,
+        tree_index: int,
+        tree_graph: "TreeGraph"
+    ) -> None:
+        """
+        Cache a TreeGraph object.
+
+        Args:
+            session_id: Session identifier
+            tree_index: Tree index in the tree sequence
+            tree_graph: TreeGraph object to cache
+        """
+        if self.redis:
+            await self._redis_set(session_id, tree_index, tree_graph)
+        else:
+            await self._local_set(session_id, tree_index, tree_graph)
+
+    async def get_all_for_session(self, session_id: str) -> Dict[int, "TreeGraph"]:
+        """
+        Retrieve all cached TreeGraphs for a session.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            Dict mapping tree_index to TreeGraph
+        """
+        if self.redis:
+            return await self._redis_get_all(session_id)
+        else:
+            return self._local_get_all(session_id)
+
+    async def clear_session(self, session_id: str) -> None:
+        """
+        Clear all cached TreeGraphs for a session.
+
+        Call this when:
+        - Session loads a new file
+        - Session is deleted/expired
+
+        Args:
+            session_id: Session identifier
+        """
+        if self.redis:
+            await self._redis_clear_session(session_id)
+        else:
+            await self._local_clear_session(session_id)
+
+    async def evict_not_visible(self, session_id: str, visible_indices: set) -> int:
+        """
+        Evict all cached TreeGraphs NOT in the visible set.
+
+        This implements visibility-based eviction: trees that are no longer
+        visible in the viewport are removed from cache.
+
+        Args:
+            session_id: Session identifier
+            visible_indices: Set of tree indices currently visible
+
+        Returns:
+            Number of trees evicted
+        """
+        if self.redis:
+            return await self._redis_evict_not_visible(session_id, visible_indices)
+        else:
+            return await self._local_evict_not_visible(session_id, visible_indices)
+
+    # ==================== Redis Implementation ====================
+
+    async def _redis_get(self, session_id: str, tree_index: int) -> Optional["TreeGraph"]:
+        """Get TreeGraph from Redis."""
+        try:
+            key = self._redis_key(session_id, tree_index)
+            data = await self.redis.get(key)
+            if data:
+                return pickle.loads(data)
+            return None
+        except Exception as e:
+            print(f"TreeGraphCache Redis get error: {e}")
+            return None
+
+    async def _redis_set(
+        self,
+        session_id: str,
+        tree_index: int,
+        tree_graph: "TreeGraph"
+    ) -> None:
+        """Set TreeGraph in Redis."""
+        try:
+            # Serialize and store
+            key = self._redis_key(session_id, tree_index)
+            data = pickle.dumps(tree_graph)
+            await self.redis.setex(key, REDIS_TTL_SECONDS, data)
+        except Exception as e:
+            print(f"TreeGraphCache Redis set error: {e}")
+
+    async def _redis_get_all(self, session_id: str) -> Dict[int, "TreeGraph"]:
+        """Get all TreeGraphs for a session from Redis."""
+        result = {}
+        try:
+            pattern = self._redis_session_pattern(session_id)
+            async for key in self.redis.scan_iter(pattern):
+                # Extract tree_index from key: treegraph:{session_id}:{tree_index}
+                parts = key.split(":")
+                if len(parts) == 3:
+                    tree_index = int(parts[2])
+                    data = await self.redis.get(key)
+                    if data:
+                        result[tree_index] = pickle.loads(data)
+        except Exception as e:
+            print(f"TreeGraphCache Redis get_all error: {e}")
+        return result
+
+    async def _redis_clear_session(self, session_id: str) -> None:
+        """Clear all TreeGraphs for a session from Redis."""
+        try:
+            pattern = self._redis_session_pattern(session_id)
+            keys = []
+            async for key in self.redis.scan_iter(pattern):
+                keys.append(key)
+            if keys:
+                await self.redis.delete(*keys)
+                print(f"TreeGraphCache cleared {len(keys)} trees for session {session_id[:8]}...")
+        except Exception as e:
+            print(f"TreeGraphCache Redis clear error: {e}")
+
+    async def _redis_evict_not_visible(self, session_id: str, visible_indices: set) -> int:
+        """Evict non-visible trees from Redis."""
+        evicted = 0
+        try:
+            pattern = self._redis_session_pattern(session_id)
+            keys_to_delete = []
+            async for key in self.redis.scan_iter(pattern):
+                # Extract tree_index from key: treegraph:{session_id}:{tree_index}
+                parts = key.decode().split(":") if isinstance(key, bytes) else key.split(":")
+                if len(parts) == 3:
+                    tree_index = int(parts[2])
+                    if tree_index not in visible_indices:
+                        keys_to_delete.append(key)
+            if keys_to_delete:
+                await self.redis.delete(*keys_to_delete)
+                evicted = len(keys_to_delete)
+                if evicted > 0:
+                    print(f"TreeGraphCache evicted {evicted} non-visible trees for session {session_id[:8]}...")
+        except Exception as e:
+            print(f"TreeGraphCache Redis evict error: {e}")
+        return evicted
+
+    # ==================== Local (In-Memory) Implementation ====================
+
+    def _local_get(self, session_id: str, tree_index: int) -> Optional["TreeGraph"]:
+        """Get TreeGraph from local cache."""
+        session_cache = self._local_cache.get(session_id)
+        if session_cache and tree_index in session_cache:
+            # Move to end for LRU (most recently used)
+            session_cache.move_to_end(tree_index)
+            return session_cache[tree_index]
+        return None
+
+    async def _local_set(
+        self,
+        session_id: str,
+        tree_index: int,
+        tree_graph: "TreeGraph"
+    ) -> None:
+        """Set TreeGraph in local cache."""
+        async with self._lock:
+            if session_id not in self._local_cache:
+                self._local_cache[session_id] = OrderedDict()
+
+            session_cache = self._local_cache[session_id]
+            session_cache[tree_index] = tree_graph
+
+    def _local_get_all(self, session_id: str) -> Dict[int, "TreeGraph"]:
+        """Get all TreeGraphs for a session from local cache."""
+        session_cache = self._local_cache.get(session_id)
+        if session_cache:
+            return dict(session_cache)
+        return {}
+
+    async def _local_clear_session(self, session_id: str) -> None:
+        """Clear all TreeGraphs for a session from local cache."""
+        async with self._lock:
+            if session_id in self._local_cache:
+                count = len(self._local_cache[session_id])
+                del self._local_cache[session_id]
+                print(f"TreeGraphCache cleared {count} trees for session {session_id[:8]}...")
+
+    async def _local_evict_not_visible(self, session_id: str, visible_indices: set) -> int:
+        """Evict non-visible trees from local cache."""
+        evicted = 0
+        async with self._lock:
+            session_cache = self._local_cache.get(session_id)
+            if session_cache:
+                keys_to_remove = [idx for idx in session_cache.keys() if idx not in visible_indices]
+                for idx in keys_to_remove:
+                    del session_cache[idx]
+                    evicted += 1
+                if evicted > 0:
+                    print(f"TreeGraphCache evicted {evicted} non-visible trees for session {session_id[:8]}...")
+        return evicted
+
+    # ==================== Utility Methods ====================
+
+    def get_stats(self) -> Dict:
+        """Get cache statistics."""
+        if self.redis:
+            return {
+                "mode": "redis",
+                "sessions": "N/A (use Redis commands to inspect)"
+            }
+        else:
+            total_trees = sum(len(cache) for cache in self._local_cache.values())
+            return {
+                "mode": "in-memory",
+                "sessions": len(self._local_cache),
+                "total_trees": total_trees,
+                "eviction_strategy": "visibility-based"
+            }