lorax-arg 0.1__py3-none-any.whl

lorax/cli.py

@@ -0,0 +1,312 @@
+"""
+Lorax CLI - Command-line interface for the Lorax backend server.
+
+Usage:
+    lorax serve --reload              # Development mode with auto-reload
+    lorax serve --host 0.0.0.0 --port 8080 --reload
+    lorax serve --gunicorn --workers 4  # Production mode with gunicorn
+    lorax config show                 # Show current configuration
+    lorax config set cache.ts_cache_size 3  # Set a config value
+"""
+import os
+import sys
+import json
+from pathlib import Path
+import click
+
+
+# Default config path for local/conda mode
+DEFAULT_CONFIG_PATH = Path.home() / ".lorax" / "config.yaml"
+
+
+@click.group()
+@click.version_option(version="1.0.0", prog_name="lorax")
+def main():
+    """Lorax Backend CLI - Tree visualization and analysis server."""
+    pass
+
+
+@main.group()
+def config():
+    """Manage Lorax configuration."""
+    pass
+
+
+@config.command("show")
+def config_show():
+    """Show current configuration."""
+    from lorax.modes import CURRENT_MODE, CURRENT_CONFIG, get_cache_dir, get_uploads_dir
+
+    click.echo(f"Mode: {CURRENT_MODE}")
+    click.echo(f"")
+    click.echo("Cache Configuration:")
+    click.echo(f"  TS Cache Size: {CURRENT_CONFIG.ts_cache_size}")
+    click.echo(f"  Config Cache Size: {CURRENT_CONFIG.config_cache_size}")
+    click.echo(f"  Metadata Cache Size: {CURRENT_CONFIG.metadata_cache_size}")
+    click.echo(f"")
+    click.echo("Disk Cache:")
+    click.echo(f"  Enabled: {CURRENT_CONFIG.disk_cache_enabled}")
+    click.echo(f"  Max Size: {CURRENT_CONFIG.disk_cache_max_gb} GB")
+    click.echo(f"  Directory: {get_cache_dir(CURRENT_CONFIG)}")
+    click.echo(f"")
+    click.echo("Connection Limits:")
+    click.echo(f"  Max Sockets Per Session: {CURRENT_CONFIG.max_sockets_per_session}")
+    click.echo(f"  Enforce Limits: {CURRENT_CONFIG.enforce_connection_limits}")
+    click.echo(f"")
+    click.echo("Directories:")
+    click.echo(f"  Data Dir: {CURRENT_CONFIG.data_dir}")
+    click.echo(f"  Uploads Dir: {get_uploads_dir(CURRENT_CONFIG)}")
+    click.echo(f"")
+    click.echo("Requirements:")
+    click.echo(f"  Redis Required: {CURRENT_CONFIG.require_redis}")
+    click.echo(f"  GCS Required: {CURRENT_CONFIG.require_gcs}")
+
+
+@config.command("init")
+@click.option("--force", is_flag=True, help="Overwrite existing config file")
+def config_init(force: bool):
+    """Initialize a config file for local mode."""
+    config_path = DEFAULT_CONFIG_PATH
+
+    if config_path.exists() and not force:
+        click.echo(f"Config file already exists: {config_path}")
+        click.echo("Use --force to overwrite")
+        return
+
+    # Create default config
+    default_config = """# Lorax Local Configuration
+# This file is used when running in local/conda mode
+
+mode: local
+
+data_dir: ~/.lorax
+
+cache:
+  ts_cache_size: 5
+  config_cache_size: 2
+  metadata_cache_size: 10
+
+disk_cache:
+  enabled: false
+  max_gb: 10
+
+server:
+  host: 127.0.0.1
+  port: 8080
+  open_browser: true
+"""
+
+    config_path.parent.mkdir(parents=True, exist_ok=True)
+    config_path.write_text(default_config)
+    click.echo(f"Created config file: {config_path}")
+
+
+@config.command("path")
+def config_path():
+    """Show config file path."""
+    click.echo(DEFAULT_CONFIG_PATH)
+
+
+@main.command("cache-status")
+def cache_status():
+    """Show cache statistics."""
+    import asyncio
+    from lorax.context import disk_cache_manager
+
+    async def get_stats():
+        return await disk_cache_manager.get_stats()
+
+    stats = asyncio.run(get_stats())
+
+    click.echo("Disk Cache Status:")
+    if not stats.get("enabled"):
+        click.echo("  Disabled")
+    else:
+        click.echo(f"  Directory: {stats.get('cache_dir')}")
+        click.echo(f"  Size: {stats.get('total_size_mb', 0)} MB / {stats.get('max_size_mb', 0)} MB ({stats.get('usage_percent', 0)}%)")
+        click.echo(f"  Files: {stats.get('file_count', 0)}")
+
+
+@main.command("cache-clear")
+@click.confirmation_option(prompt="Are you sure you want to clear the disk cache?")
+def cache_clear():
+    """Clear the disk cache."""
+    import asyncio
+    from lorax.context import disk_cache_manager
+
+    async def clear():
+        await disk_cache_manager.clear()
+
+    asyncio.run(clear())
+    click.echo("Disk cache cleared")
+
+
+@main.command()
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    help="Host to bind to (default: 127.0.0.1 for local, 0.0.0.0 for production)"
+)
+@click.option(
+    "--port",
+    default=8080,
+    type=int,
+    help="Port to bind to (default: 8080)"
+)
+@click.option(
+    "--reload",
+    is_flag=True,
+    help="Enable auto-reload for development"
+)
+@click.option(
+    "--gunicorn",
+    is_flag=True,
+    help="Use gunicorn for production (requires lorax[prod])"
+)
+@click.option(
+    "--workers",
+    default=4,
+    type=int,
+    help="Number of gunicorn workers (default: 4, only used with --gunicorn)"
+)
+@click.option(
+    "--config",
+    default=None,
+    type=click.Path(exists=True),
+    help="Path to gunicorn config file (optional)"
+)
+@click.option(
+    "--open-browser",
+    is_flag=True,
+    help="Open browser automatically (for local mode)"
+)
+def serve(host: str, port: int, reload: bool, gunicorn: bool, workers: int, config: str, open_browser: bool):
+    """Start the Lorax backend server.
+
+    Examples:
+        # Local mode (opens browser)
+        lorax serve --open-browser
+
+        # Development mode (with auto-reload)
+        lorax serve --reload
+
+        # Specify host/port
+        lorax serve --host 0.0.0.0 --port 8080 --reload
+
+        # Production mode (with gunicorn)
+        lorax serve --gunicorn --workers 4
+    """
+    if gunicorn:
+        _run_with_gunicorn(host, port, workers, config)
+    else:
+        _run_with_uvicorn(host, port, reload, open_browser)
+
+
+def _run_with_uvicorn(host: str, port: int, reload: bool, open_browser: bool = False):
+    """Run the server with uvicorn (development/local mode)."""
+    import uvicorn
+
+    click.echo(f"Starting Lorax server with uvicorn on {host}:{port}")
+    if reload:
+        click.echo("Auto-reload enabled")
+
+    # Open browser if requested (in a separate thread to not block server start)
+    if open_browser:
+        import threading
+        import webbrowser
+        import time
+
+        def open_browser_delayed():
+            time.sleep(1.5)  # Wait for server to start
+            url = f"http://{host}:{port}" if host != "0.0.0.0" else f"http://127.0.0.1:{port}"
+            click.echo(f"Opening browser: {url}")
+            webbrowser.open(url)
+
+        threading.Thread(target=open_browser_delayed, daemon=True).start()
+
+    uvicorn.run(
+        "lorax.lorax_app:sio_app",
+        host=host,
+        port=port,
+        reload=reload,
+        log_level="info"
+    )
+
+
+def _run_with_gunicorn(host: str, port: int, workers: int, config: str):
+    """Run the server with gunicorn (production mode)."""
+    try:
+        import gunicorn.app.base
+    except ImportError:
+        click.echo(
+            "Error: gunicorn is not installed. "
+            "Install with: pip install lorax[prod]",
+            err=True
+        )
+        sys.exit(1)
+    
+    from gunicorn.app.base import BaseApplication
+    
+    class LoraxApplication(BaseApplication):
+        def __init__(self, app_uri: str, options: dict = None):
+            self.app_uri = app_uri
+            self.options = options or {}
+            super().__init__()
+        
+        def load_config(self):
+            for key, value in self.options.items():
+                if key in self.cfg.settings and value is not None:
+                    self.cfg.set(key.lower(), value)
+        
+        def load(self):
+            # Return the app URI string for gunicorn to import
+            return None
+    
+    bind = f"{host}:{port}"
+    
+    options = {
+        "bind": bind,
+        "workers": workers,
+        "worker_class": "uvicorn.workers.UvicornWorker",
+        "timeout": 0,
+        "graceful_timeout": 0,
+        "keepalive": 0,
+        "accesslog": "-",
+        "errorlog": "-",
+        "loglevel": "info",
+    }
+    
+    # Load custom config if provided
+    if config:
+        click.echo(f"Loading gunicorn config from: {config}")
+        # The config file will be loaded by gunicorn
+        options["config"] = config
+    
+    click.echo(f"Starting Lorax server with gunicorn on {bind}")
+    click.echo(f"Workers: {workers}")
+    
+    # Use subprocess to run gunicorn with the app
+    import subprocess
+    
+    cmd = [
+        sys.executable, "-m", "gunicorn",
+        "-b", bind,
+        "-w", str(workers),
+        "-k", "uvicorn.workers.UvicornWorker",
+        "--timeout", "0",
+        "--graceful-timeout", "0",
+        "--access-logfile", "-",
+        "--error-logfile", "-",
+        "--log-level", "info",
+    ]
+    
+    if config:
+        cmd.extend(["-c", config])
+    
+    cmd.append("lorax.lorax_app:sio_app")
+    
+    subprocess.run(cmd)
+
+
+if __name__ == "__main__":
+    main()

lorax/cloud/gcs_utils.py

@@ -0,0 +1,205 @@
+from google.cloud import storage
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+import zoneinfo
+import aiofiles
+import asyncio
+import requests
+import aiohttp
+import os
+
+BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
+
+def get_gcs_client():
+    return storage.Client()
+
+
+async def _download_gcs_file_direct(bucket_name: str, blob_path: str, local_path: str):
+    """
+    Internal: Download a public GCS blob to a local file.
+
+    Args:
+        bucket_name (str): Name of the public GCS bucket (e.g. "lorax_projects")
+        blob_path (str): Path inside the bucket (e.g. "1000Genomes/1kg_chr20.trees.tsz")
+        local_path (str): Local destination path (e.g. "uploads/1kg_chr20.trees.tsz")
+    """
+    url = f"https://storage.googleapis.com/{bucket_name}/{blob_path}"
+    Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url) as resp:
+            if resp.status != 200:
+                raise RuntimeError(f"Failed to download {url} (HTTP {resp.status})")
+
+            async with aiofiles.open(local_path, "wb") as f:
+                while chunk := await resp.content.read(1024 * 1024):  # 1 MB chunks
+                    await f.write(chunk)
+
+    return Path(local_path)
+
+
+async def download_gcs_file(bucket_name: str, blob_path: str, local_path: str):
+    """
+    Asynchronously download a public GCS blob to a local file.
+    Uses disk cache if available for automatic caching and eviction.
+
+    Args:
+        bucket_name (str): Name of the public GCS bucket (e.g. "lorax_projects")
+        blob_path (str): Path inside the bucket (e.g. "1000Genomes/1kg_chr20.trees.tsz")
+        local_path (str): Local destination path (e.g. "uploads/1kg_chr20.trees.tsz")
+
+    Returns:
+        Path to the downloaded file
+    """
+    # Import here to avoid circular imports
+    from lorax.context import disk_cache_manager
+
+    if disk_cache_manager.enabled:
+        # Use disk cache for managed downloads
+        async def download_func(path: str):
+            await _download_gcs_file_direct(bucket_name, blob_path, path)
+
+        cached_path = await disk_cache_manager.get_or_download(
+            bucket_name, blob_path, download_func
+        )
+
+        # If local_path differs from cached path, create a symlink or copy
+        local_path_obj = Path(local_path)
+        if cached_path != local_path_obj:
+            local_path_obj.parent.mkdir(parents=True, exist_ok=True)
+            # Symlink to cached file (avoid duplicate storage)
+            if local_path_obj.exists() or local_path_obj.is_symlink():
+                local_path_obj.unlink()
+            try:
+                local_path_obj.symlink_to(cached_path)
+            except OSError:
+                # Fallback: just return cached path if symlink fails
+                return cached_path
+
+        return local_path_obj
+    else:
+        # Direct download without caching
+        return await _download_gcs_file_direct(bucket_name, blob_path, local_path)
+
+
+async def download_gcs_file_cached(
+    bucket_name: str,
+    blob_path: str,
+    disk_cache_manager=None
+) -> Optional[Path]:
+    """
+    Download a GCS file using the disk cache manager.
+
+    This is the preferred method for production use - it handles:
+    - Distributed locking (prevents duplicate downloads across workers)
+    - LRU eviction (manages disk space)
+    - Access time tracking
+
+    Args:
+        bucket_name: GCS bucket name
+        blob_path: Path within bucket
+        disk_cache_manager: Optional cache manager (uses global if not provided)
+
+    Returns:
+        Path to cached file, or None if download failed
+    """
+    if disk_cache_manager is None:
+        from lorax.context import disk_cache_manager
+
+    if not disk_cache_manager.enabled:
+        return None
+
+    async def download_func(local_path: str):
+        await _download_gcs_file_direct(bucket_name, blob_path, local_path)
+
+    try:
+        return await disk_cache_manager.get_or_download(
+            bucket_name, blob_path, download_func
+        )
+    except Exception as e:
+        print(f"Failed to download {blob_path} from GCS: {e}")
+        return None
+
+def get_public_gcs_dict(
+    bucket_name: str,
+    sid: str,
+    prefix: str = "",
+    projects=None,
+    include_uploads: bool = False,
+    uploads_sid: str | None = None,
+):
+    if projects is None:
+        projects = {}
+    api_url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o"
+    params = {"prefix": '', "fields": "items(name)"}
+    resp = requests.get(api_url, params=params)
+    resp.raise_for_status()
+    items = resp.json().get("items", [])
+
+    for item in items:
+        name = item['name']
+        path_parts = name.split("/")
+        
+        # Must have at least a top-level directory (e.g., 'folder/')
+        if len(path_parts) < 2:
+            continue
+        name_first = path_parts[0]
+        second_part = path_parts[1]
+
+        # Handle Uploads filtering
+        if name_first == 'Uploads':
+            if not include_uploads:
+                continue
+            if not uploads_sid or second_part != uploads_sid:
+                continue
+            # Require a filename component
+            if len(path_parts) < 3:
+                continue
+
+        if name_first not in projects:
+            projects[name_first] = {'folder': name_first,'files': [], 'description': ''}
+
+        if name_first != 'Uploads' and second_part:
+            if second_part not in projects[name_first]['files']:
+                projects[name_first]['files'].append(second_part)
+        elif name_first == 'Uploads':
+            filename = path_parts[2]
+            if filename not in projects[name_first]['files']:
+                projects[name_first]['files'].append(filename)
+    return projects
+    
+async def upload_to_gcs(bucket_name: str, local_path: Path, sid: str):
+    """
+    Upload file to GCS under session-specific folder.
+
+    Args:
+        bucket_name (str): Target GCS bucket
+        local_path (Path): Local file to upload
+        sid (str): Session ID (used as prefix)
+    """
+    blob_path = f"Uploads/{sid}/{local_path.name}"
+
+    # Use executor to avoid blocking event loop
+    loop = asyncio.get_event_loop()
+    await loop.run_in_executor(None, _upload_file_sync, bucket_name, local_path, blob_path)
+    return f"https://storage.googleapis.com/{bucket_name}/{blob_path}"
+
+def _upload_file_sync(bucket_name: str, local_path: Path, blob_path: str):
+    """Synchronous helper called within a thread executor."""
+    client = get_gcs_client()
+    bucket = client.bucket(bucket_name)
+    blob = bucket.blob(blob_path)
+    blob.upload_from_filename(str(local_path))
+    blob.custom_time = datetime.utcnow()
+    pacific_tz = zoneinfo.ZoneInfo("America/Los_Angeles")
+    now_pacific = datetime.now(pacific_tz)
+
+    # Convert to UTC (GCS expects UTC timestamps)
+    now_utc = now_pacific.astimezone(timezone.utc)
+
+    # Assign custom_time so GCS lifecycle rule can delete it after 7 days
+    blob.custom_time = now_utc
+    blob.patch()
+    # blob.make_public()  # Optional: make public for browser access
+    print(f"Uploaded {local_path.name} to gs://{bucket_name}/{blob_path}")

lorax/constants.py

@@ -0,0 +1,66 @@
+"""
+Lorax Backend Constants
+
+Centralized configuration values to avoid hardcoding throughout the codebase.
+Mode-aware configuration based on deployment environment.
+"""
+
+import os
+from pathlib import Path
+
+# Import mode configuration
+from lorax.modes import (
+    get_mode_config,
+    get_uploads_dir,
+    get_cache_dir,
+    CURRENT_MODE,
+    CURRENT_CONFIG,
+)
+
+# Session Configuration
+SESSION_COOKIE = "lorax_sid"
+COOKIE_MAX_AGE = 7 * 24 * 60 * 60  # 7 days in seconds
+
+# Cache Configuration (mode-aware)
+TS_CACHE_SIZE = CURRENT_CONFIG.ts_cache_size
+CONFIG_CACHE_SIZE = CURRENT_CONFIG.config_cache_size
+METADATA_CACHE_SIZE = CURRENT_CONFIG.metadata_cache_size
+
+# Disk Cache Configuration (mode-aware)
+DISK_CACHE_ENABLED = CURRENT_CONFIG.disk_cache_enabled
+DISK_CACHE_DIR = get_cache_dir(CURRENT_CONFIG)
+DISK_CACHE_MAX_BYTES = CURRENT_CONFIG.disk_cache_max_gb * 1024 * 1024 * 1024
+
+# Connection Limits (mode-aware)
+MAX_SOCKETS_PER_SESSION = CURRENT_CONFIG.max_sockets_per_session
+ENFORCE_CONNECTION_LIMITS = CURRENT_CONFIG.enforce_connection_limits
+
+# File Types
+SUPPORTED_EXTENSIONS = {'.tsz', '.trees', '.csv'}
+
+# Directory Names (mode-aware)
+UPLOADS_DIR = str(get_uploads_dir(CURRENT_CONFIG))
+
+# Default Values
+DEFAULT_WINDOW_SIZE = 50000
+
+# Socket.IO Configuration
+SOCKET_PING_TIMEOUT = 60  # seconds
+SOCKET_PING_INTERVAL = 25  # seconds
+MAX_HTTP_BUFFER_SIZE = 50_000_000  # 50 MB
+
+# Error Codes
+ERROR_SESSION_NOT_FOUND = "SESSION_NOT_FOUND"
+ERROR_MISSING_SESSION = "MISSING_SESSION"
+ERROR_NO_FILE_LOADED = "NO_FILE_LOADED"
+ERROR_TOO_MANY_CONNECTIONS = "TOO_MANY_CONNECTIONS"
+ERROR_CONNECTION_REPLACED = "CONNECTION_REPLACED"
+
+
+def print_config():
+    """Print current configuration for debugging."""
+    print(f"Mode: {CURRENT_MODE}")
+    print(f"TS Cache Size: {TS_CACHE_SIZE}")
+    print(f"Disk Cache: {DISK_CACHE_ENABLED} ({CURRENT_CONFIG.disk_cache_max_gb}GB)")
+    print(f"Max Sockets/Session: {MAX_SOCKETS_PER_SESSION}")
+    print(f"Uploads Dir: {UPLOADS_DIR}")

lorax/context.py

@@ -0,0 +1,80 @@
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Import mode configuration first (before other lorax imports that depend on it)
+from lorax.modes import (
+    CURRENT_MODE,
+    CURRENT_CONFIG,
+    get_cache_dir,
+    validate_mode_requirements,
+)
+from lorax.session_manager import SessionManager
+from lorax.redis_utils import create_redis_client, get_redis_config
+from lorax.cache import DiskCacheManager, TreeGraphCache, CsvTreeGraphCache
+from lorax.constants import (
+    DISK_CACHE_ENABLED,
+    DISK_CACHE_DIR,
+    DISK_CACHE_MAX_BYTES,
+)
+
+# Validate mode requirements
+validation_errors = validate_mode_requirements(CURRENT_CONFIG)
+if validation_errors:
+    for error in validation_errors:
+        print(f"Warning: {error}")
+
+# Shared Global State
+# We initialize these singletons here to ensure all modules share the same instances
+# This is critical for in-memory mode so routes and sockets share the same stores
+
+REDIS_CLUSTER_URL, REDIS_CLUSTER = get_redis_config()
+session_manager = SessionManager(redis_url=REDIS_CLUSTER_URL, redis_cluster=REDIS_CLUSTER)
+
+# Common Environment Variables
+BUCKET_NAME = os.getenv("BUCKET_NAME") or os.getenv("GCS_BUCKET_NAME") or 'lorax_projects'
+
+# Initialize Disk Cache Manager
+# Uses Redis for distributed locking if available, falls back to file locks
+_redis_client = None
+if REDIS_CLUSTER_URL and DISK_CACHE_ENABLED:
+    try:
+        _redis_client = create_redis_client(
+            REDIS_CLUSTER_URL,
+            decode_responses=True,
+            cluster=REDIS_CLUSTER,
+        )
+        print("DiskCacheManager using Redis for distributed locking")
+    except Exception as e:
+        print(f"Warning: Failed to connect Redis for disk cache: {e}")
+
+disk_cache_manager = DiskCacheManager(
+    cache_dir=Path(DISK_CACHE_DIR),
+    max_size_bytes=DISK_CACHE_MAX_BYTES,
+    redis_client=_redis_client,
+    enabled=DISK_CACHE_ENABLED,
+)
+
+# Initialize TreeGraph Cache for per-session tree caching
+# Uses Redis in production for distributed caching, in-memory for local mode
+_tree_graph_redis = None
+if REDIS_CLUSTER_URL:
+    try:
+        # Create a separate connection for binary data (decode_responses=False)
+        _tree_graph_redis = create_redis_client(
+            REDIS_CLUSTER_URL,
+            decode_responses=False,
+            cluster=REDIS_CLUSTER,
+        )
+        print(f"TreeGraphCache using Redis at {REDIS_CLUSTER_URL}")
+    except Exception as e:
+        print(f"Warning: Failed to connect Redis for TreeGraphCache: {e}")
+
+tree_graph_cache = TreeGraphCache(redis_client=None)
+
+# CSV mode: cache parsed Newick trees per session (in-memory only)
+csv_tree_graph_cache = CsvTreeGraphCache()
+
+print(f"Context initialized: mode={CURRENT_MODE}, disk_cache={DISK_CACHE_ENABLED}")

lorax/csv/__init__.py

@@ -0,0 +1,7 @@
+"""CSV support for Lorax backend.
+
+This package is intentionally minimal: it provides a config builder that matches
+the existing frontend contract and a layout builder that can return an empty but
+valid PyArrow IPC buffer for early CSV support.
+"""
+