lorax-arg 0.1__py3-none-any.whl

lorax/modes.py

@@ -0,0 +1,190 @@
+"""
+Lorax Deployment Mode Detection and Configuration
+
+Supports three deployment modes:
+- local: Conda package / desktop use (single user, no Redis/GCS)
+- development: Developer mode (single process, optional GCS)
+- production: Cloud deployment (multi-worker, Redis + GCS required)
+"""
+
+import os
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class ModeConfig:
+    """Configuration for a deployment mode."""
+    mode: str
+    ts_cache_size: int
+    config_cache_size: int
+    metadata_cache_size: int
+    disk_cache_enabled: bool
+    disk_cache_max_gb: int
+    max_sockets_per_session: int
+    enforce_connection_limits: bool
+    data_dir: Path
+    require_redis: bool
+    require_gcs: bool
+
+
+# Mode-specific defaults
+MODE_CONFIGS = {
+    "local": ModeConfig(
+        mode="local",
+        ts_cache_size=5,
+        config_cache_size=2,
+        metadata_cache_size=10,
+        disk_cache_enabled=False,  # No GCS in local mode
+        disk_cache_max_gb=10,
+        max_sockets_per_session=100,  # Effectively unlimited
+        enforce_connection_limits=False,
+        data_dir=Path.home() / ".lorax",
+        require_redis=False,
+        require_gcs=False,
+    ),
+    "development": ModeConfig(
+        mode="development",
+        ts_cache_size=5,
+        config_cache_size=2,
+        metadata_cache_size=10,
+        disk_cache_enabled=True,
+        disk_cache_max_gb=10,
+        max_sockets_per_session=100,  # Relaxed for testing
+        enforce_connection_limits=False,
+        data_dir=Path("UPLOADS"),
+        require_redis=False,
+        require_gcs=False,
+    ),
+    "production": ModeConfig(
+        mode="production",
+        ts_cache_size=2,  # Lower for memory efficiency with many workers
+        config_cache_size=2,
+        metadata_cache_size=5,
+        disk_cache_enabled=True,
+        disk_cache_max_gb=50,
+        max_sockets_per_session=5,
+        enforce_connection_limits=True,
+        data_dir=Path("/var/lorax"),
+        require_redis=True,
+        require_gcs=True,
+    ),
+}
+
+
+def detect_mode() -> str:
+    """
+    Detect the deployment mode based on environment variables.
+
+    Priority:
+    1. Explicit LORAX_MODE environment variable
+    2. Auto-detect based on REDIS_CLUSTER and GCS_BUCKET_NAME
+    3. Default to 'local' for conda package usage
+    """
+    explicit_mode = os.getenv("LORAX_MODE", "").lower()
+    if explicit_mode in MODE_CONFIGS:
+        return explicit_mode
+
+    # Auto-detect based on environment
+    has_redis = bool(os.getenv("REDIS_CLUSTER"))
+    has_gcs = bool(os.getenv("GCS_BUCKET_NAME") or os.getenv("BUCKET_NAME"))
+    if has_redis and has_gcs:
+        return "production"
+    elif has_gcs:
+        return "development"
+    else:
+        return "local"
+
+
+def get_mode_config(mode: Optional[str] = None) -> ModeConfig:
+    """
+    Get configuration for a specific mode.
+
+    Args:
+        mode: Mode name or None to auto-detect
+
+    Returns:
+        ModeConfig with all settings for the mode
+    """
+    if mode is None:
+        mode = detect_mode()
+
+    config = MODE_CONFIGS.get(mode, MODE_CONFIGS["local"])
+
+    # Apply environment overrides
+    return ModeConfig(
+        mode=config.mode,
+        ts_cache_size=int(os.getenv("TS_CACHE_SIZE", config.ts_cache_size)),
+        config_cache_size=int(os.getenv("CONFIG_CACHE_SIZE", config.config_cache_size)),
+        metadata_cache_size=int(os.getenv("METADATA_CACHE_SIZE", config.metadata_cache_size)),
+        disk_cache_enabled=config.disk_cache_enabled,
+        disk_cache_max_gb=int(os.getenv("DISK_CACHE_MAX_GB", config.disk_cache_max_gb)),
+        max_sockets_per_session=int(os.getenv("MAX_SOCKETS_PER_SESSION", config.max_sockets_per_session)),
+        enforce_connection_limits=config.enforce_connection_limits,
+        data_dir=Path(os.getenv("LORAX_DATA_DIR", config.data_dir)),
+        require_redis=config.require_redis,
+        require_gcs=config.require_gcs,
+    )
+
+
+def get_data_dir(mode_config: Optional[ModeConfig] = None) -> Path:
+    """Get the data directory for the current mode."""
+    if mode_config is None:
+        mode_config = get_mode_config()
+    mode_config.data_dir.mkdir(parents=True, exist_ok=True)
+    return mode_config.data_dir
+
+
+def get_uploads_dir(mode_config: Optional[ModeConfig] = None) -> Path:
+    """Get the uploads directory for the current mode."""
+    data_dir = get_data_dir(mode_config)
+    uploads_dir = data_dir / "uploads" if mode_config and mode_config.mode == "local" else data_dir
+    uploads_dir.mkdir(parents=True, exist_ok=True)
+    return uploads_dir
+
+
+def get_cache_dir(mode_config: Optional[ModeConfig] = None) -> Path:
+    """Get the disk cache directory for the current mode."""
+    if mode_config is None:
+        mode_config = get_mode_config()
+
+    cache_dir_env = os.getenv("DISK_CACHE_DIR")
+    if cache_dir_env:
+        cache_dir = Path(cache_dir_env)
+    elif mode_config.mode == "local":
+        cache_dir = mode_config.data_dir / "cache"
+    elif mode_config.mode == "development":
+        cache_dir = Path("/tmp/lorax_cache")
+    else:
+        cache_dir = Path("/var/lorax/cache")
+
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def validate_mode_requirements(mode_config: ModeConfig) -> list:
+    """
+    Validate that required services are available for the mode.
+
+    Returns:
+        List of validation errors (empty if all requirements met)
+    """
+    errors = []
+
+    if mode_config.require_redis and not os.getenv("REDIS_CLUSTER"):
+        errors.append(f"{mode_config.mode} mode requires REDIS_CLUSTER environment variable")
+
+    if mode_config.require_gcs:
+        bucket = os.getenv("GCS_BUCKET_NAME") or os.getenv("BUCKET_NAME")
+        if not bucket:
+            errors.append(f"{mode_config.mode} mode requires GCS_BUCKET_NAME environment variable")
+
+    return errors
+
+
+# Initialize on import
+CURRENT_MODE = detect_mode()
+CURRENT_CONFIG = get_mode_config(CURRENT_MODE)
+
+print(f"Lorax mode: {CURRENT_MODE}")

lorax/pg.py

@@ -0,0 +1,183 @@
+
+import argparse
+import sys
+import math
+
+try:
+    import tskit  # noqa: F401
+except Exception as e:
+    # We allow the script to be inspected without tskit, but bail at runtime.
+    pass
+
+
+def clamp(v, lo, hi):
+    return max(lo, min(hi, v))
+
+
+def write_track_line(out, name, description, visibility=2, use_score=True, item_rgb=True):
+    parts = [f'track name="{name}"']
+    if description:
+        parts.append(f'description="{description}"')
+    parts.append(f'visibility={int(visibility)}')
+    if use_score:
+        parts.append("useScore=1")
+    if item_rgb:
+        parts.append('itemRgb="On"')
+    out.write(" ".join(parts) + "\n")
+
+
+def tree_stat(tree, stat):
+    """
+    Return a per-tree statistic used for score/color.
+    Currently supported:
+      - root_time: time of the (single) root; if multiple roots, max root time
+      - total_branch_length: sum of branch lengths (in time units)
+      - n_roots: number of roots (integer)
+    """
+    if stat == "root_time":
+        # If there's a single root, use that; if multiple, take the max time.
+        if hasattr(tree, "roots"):
+            roots = list(tree.roots)
+        else:
+            # older tskit: use tree.root
+            roots = [tree.root] if tree.root != tskit.NULL else []
+        if not roots:
+            return 0.0
+        return max(tree.time(r) for r in roots)
+
+    if stat == "total_branch_length":
+        # Sum over edges in this tree: sum_t (time(parent) - time(child))
+        total = 0.0
+        for u in tree.nodes():
+            p = tree.parent(u)
+            if p != tskit.NULL:
+                total += tree.time(p) - tree.time(u)
+        return max(0.0, total)
+
+    if stat == "n_roots":
+        return float(len(list(tree.roots)))
+
+    raise ValueError(f"Unsupported stat: {stat}")
+
+
+def main():
+    p = argparse.ArgumentParser(
+        description="Create a BED file from a tskit tree sequence."
+    )
+    p.add_argument("input", help="Input .trees (tskit tree sequence)")
+    p.add_argument("output", help="Output .bed path, or '-' for stdout")
+    p.add_argument("--chrom", default="chr1", help="Chromosome/contig name for BED (default: chr1)")
+    p.add_argument("--mode", choices=["trees", "sites"], default="trees",
+                   help="Emit intervals per tree or per site (default: trees)")
+    p.add_argument("--stat", choices=["root_time", "total_branch_length", "n_roots"],
+                   default="root_time", help="Per-tree stat for score/color (trees mode)")
+    p.add_argument("--ucsc-track", action="store_true",
+                   help="Prepend a UCSC track line to the BED output")
+    p.add_argument("--track-name", default="tskit-BED",
+                   help="Name to use in UCSC track line (if --ucsc-track)")
+    p.add_argument("--track-description", default="Generated by tskit_to_bed.py",
+                   help="Description to use in UCSC track line")
+    p.add_argument("--no-bed9", action="store_true",
+                   help="Write simple 3-column BED (chrom, start, end). Default is 9-column BED with score/color.")
+    p.add_argument("--precision", type=int, default=0,
+                   help="Round interval coordinates to this many decimals before casting to int (default: 0)")
+
+    args = p.parse_args()
+
+    ts = tskit.load(args.input)
+
+    # Choose output
+    if args.output == "-":
+        out = sys.stdout
+        close_out = False
+    else:
+        out = open(args.output, "w", encoding="utf-8")
+        close_out = True
+
+    try:
+        if args.ucsc_track:
+            # continue
+            write_track_line(
+                out,
+                name=args.track_name,
+                description=args.track_description,
+                visibility=2,
+                use_score=not args.no_bed9,
+                item_rgb=not args.no_bed9,
+            )
+
+        chrom = args.chrom
+
+        if args.mode == "sites":
+            # Emit per-site 3-col or 6/9-col BED. We keep it simple: 3-col unless bed9 is desired.
+            for var in ts.variants():
+                # tskit positions are float; BED wants ints. Round then cast.
+                start = int(round(var.position, args.precision))
+                end = start + 1
+                if args.no_bed9:
+                    out.write(f"{chrom}\t{start}\t{end}\n")
+                else:
+                    name = getattr(var, "id", None)
+                    if name is None:
+                        # fallback to pos + first alt allele
+                        alleles = [a for a in var.alleles if a is not None]
+                        name = f"pos{start}:{'/'.join(alleles[:2])}" if alleles else f"pos{start}"
+                    score = 0
+                    strand = "."
+                    thickStart = start
+                    thickEnd = end
+                    rgb = "0,0,0"
+                    out.write(f"{chrom}\t{start}\t{end}\t{name}\t{score}\t{strand}\t{thickStart}\t{thickEnd}\t{rgb}\n")
+            return
+
+        # trees mode
+        # First pass: collect stats and intervals
+        intervals = []
+        stat_values = []
+        for tree in ts.trees(tracked_samples=None):
+            left = int(round(tree.interval.left, args.precision))
+            right = int(round(tree.interval.right, args.precision))
+            if right <= left:
+                continue
+            val = tree_stat(tree, args.stat)
+            intervals.append((left, right, tree.index, val))
+            stat_values.append(val)
+
+        if not intervals:
+            return
+
+        # Determine scaling for score/color
+        vmax = max(stat_values) if stat_values else 1.0
+        if vmax <= 0:
+            vmax = 1.0
+
+        for left, right, tidx, val in intervals:
+            if args.no_bed9:
+                out.write(f"{chrom}\t{left}\t{right}\n")
+                continue
+
+            # name, score, strand, thickStart, thickEnd, itemRgb
+            name = f"tree_{tidx}"
+            score = int(round(1000.0 * (val / vmax)))
+            score = clamp(score, 0, 1000)
+            strand = "."
+            thickStart = left
+            thickEnd = right
+
+            # Map val to grayscale 0..255
+            if vmax > 0:
+                g = int(round(255.0 * (val / vmax)))
+            else:
+                g = 0
+            g = clamp(g, 0, 255)
+            rgb = f"{g},{g},{g}"
+
+            out.write(f"{chrom}\t{left}\t{right}\t{name}\t{score}\t{strand}\t{thickStart}\t{thickEnd}\t{rgb}\n")
+
+    finally:
+        if close_out:
+            out.close()
+
+
+if __name__ == "__main__":
+    main()

lorax/redis_utils.py

@@ -0,0 +1,30 @@
+import os
+from typing import Optional, Tuple
+
+
+def get_redis_config() -> Tuple[Optional[str], bool]:
+    """
+    Returns (redis_url, is_cluster).
+
+    - REDIS_CLUSTER provides the cluster endpoint URL.
+    """
+    url = os.getenv("REDIS_CLUSTER")
+    is_cluster = bool(url)
+    return url, is_cluster
+
+
+def create_redis_client(
+    redis_url: Optional[str],
+    *,
+    decode_responses: bool,
+    cluster: bool,
+):
+    if not redis_url:
+        return None
+
+    if cluster:
+        from redis.asyncio.cluster import RedisCluster
+        return RedisCluster.from_url(redis_url, decode_responses=decode_responses)
+
+    import redis.asyncio as aioredis
+    return aioredis.from_url(redis_url, decode_responses=decode_responses)

lorax/routes.py

@@ -0,0 +1,137 @@
+import os
+import json
+from pathlib import Path
+from typing import Optional
+import aiofiles
+
+from fastapi import APIRouter, Request, Response, UploadFile, File, Query
+from fastapi.responses import JSONResponse
+
+from lorax.context import session_manager, BUCKET_NAME
+from lorax.modes import CURRENT_MODE
+from lorax.constants import UPLOADS_DIR
+from lorax.cloud.gcs_utils import upload_to_gcs
+from lorax.handlers import (
+    handle_upload,
+    get_projects,
+    cache_status,
+)
+
+router = APIRouter()
+UPLOAD_DIR = Path(UPLOADS_DIR)
+UPLOAD_DIR.mkdir(exist_ok=True)
+
+@router.get("/health")
+async def healthz():
+    redis_ok = await session_manager.health_check()
+    return {"ok": True, "redis": redis_ok}
+
+@router.get("/fevicon.ico")
+async def favicon():
+    return Response(content="", media_type="image/x-icon")
+
+@router.get("/")
+async def root():
+    return Response(content="Lorax Backend is running...", media_type="text/html")
+
+@router.post("/init-session")
+async def init_session(request: Request, response: Response):
+    sid, session = await session_manager.get_or_create_session(request, response)
+    print("init-session:", sid)
+    return {"sid": sid}
+
+@router.get("/projects")
+async def projects(request: Request, response: Response):
+    sid, session = await session_manager.get_or_create_session(request, response)
+    projects = await get_projects(UPLOAD_DIR, BUCKET_NAME, sid=sid)
+    return {"projects": projects}
+
+@router.get("/memory_status")
+async def memory_status():
+    print("cache-status")
+    return await cache_status()
+
+@router.get("/{file}")
+async def get_file(
+    request: Request,
+    response: Response,
+    file: Optional[str] = None,
+    project: Optional[str] = Query(None),
+    genomiccoordstart: Optional[int] = Query(None),
+    genomiccoordend: Optional[int] = Query(None),
+    share_sid: Optional[str] = Query(None),
+):
+    sid, session = await session_manager.get_or_create_session(request, response)
+    if project == "Uploads" and share_sid and share_sid != sid:
+        print(f"⚠️ share_sid denied for sid={sid} target={share_sid}")
+        return JSONResponse(
+            status_code=403,
+            content={"error": "Access denied for shared upload."},
+        )
+    if file and file != "" and file != "ucgb":
+        if project == 'Uploads':
+            target_sid = share_sid if share_sid else sid
+            file_path = UPLOAD_DIR / project / target_sid / file
+        else:
+            file_path = UPLOAD_DIR / (project or "") / file
+    else:
+        file = "1kg_chr20.trees.tsz"
+        file_path = UPLOAD_DIR / (project or "1000Genomes") / file
+    try:
+        ctx = await handle_upload(str(file_path), str(UPLOAD_DIR))
+        viz_config = ctx.config
+
+        # Override initial_position if client provided genomic coordinates
+        if genomiccoordstart is not None and genomiccoordend is not None:
+            viz_config['initial_position'] = [genomiccoordstart, genomiccoordend]
+
+        # Update session with loaded file path
+        session.file_path = str(file_path)
+        await session_manager.save_session(session)
+
+    except Exception as e:
+        print(f"❌ Error loading file: {e}")
+        return {"error": str(e)}
+
+    return {
+        "sid": sid,
+        "file": file,
+        "project": project,
+        "config": viz_config,
+        "genomiccoordstart": genomiccoordstart,
+        "genomiccoordend": genomiccoordend,
+    }
+
+@router.post("/upload")
+async def upload(request: Request, response: Response, file: UploadFile = File(...)):
+    """
+    Upload a file to the server. Stores locally for all modes; uploads to GCS
+    only when not running in local mode.
+    """
+    sid, session = await session_manager.get_or_create_session(request, response)
+
+    # Local mode: flat under Uploads; Non-local: session-scoped folder
+    if CURRENT_MODE == "local":
+        user_dir = UPLOAD_DIR / "Uploads"
+    else:
+        user_dir = UPLOAD_DIR / "Uploads" / sid
+    user_dir.mkdir(parents=True, exist_ok=True)
+
+    file_path = user_dir / file.filename
+
+    try:
+        async with aiofiles.open(file_path, "wb") as f:
+            while chunk := await file.read(1024 * 1024):
+                await f.write(chunk)
+
+        # Upload to GCS asynchronously when allowed
+        if CURRENT_MODE != "local" and BUCKET_NAME:
+            await upload_to_gcs(BUCKET_NAME, file_path, sid)
+        
+        return JSONResponse(
+            status_code=200,
+            content={"message": "File uploaded", "sid": sid, "owner_sid": sid, "filename": file.filename}
+        )
+    except Exception as e:
+        print("❌ Upload error:", e)
+        return JSONResponse(status_code=500, content={"error": "Upload error"})