Sync main.py from blood-brain-omics repo

main.py

@@ -33,7 +33,6 @@ import logging
 import os
 import time
 import traceback
-import urllib.request
 from typing import Optional
 
 import duckdb
@@ -51,13 +50,10 @@ logger = logging.getLogger("blood_brain_api")
 # Configuration
 # =============================================================================
 
-# HF Dataset repo for remote Parquet queries
-HF_DATASET = os.environ.get(
-    "HF_DATASET", "stasaking/blood-brain-benchmark"
-)
-HF_BASE_URL = f"https://huggingface.co/datasets/{HF_DATASET}/resolve/main"
-
-# Local fallback for dev
+# Local data directory. Holds parquets + registry. On the HF Space the
+# start.sh entrypoint downloads these from the HF bucket
+# (stasaking/blood-brain-benchmark) into /app/data before launching the
+# API. In dev this points at webapp/data with the files already on disk.
 DATA_DIR = os.environ.get(
     "BENCHMARK_DATA_DIR",
     os.path.join(os.path.dirname(__file__), "..", "data")
@@ -219,14 +215,6 @@ db = None
 registry = None
 
 
-def _resolve_parquet(filename: str) -> str:
-    """Return remote URL or local path for a Parquet file."""
-    local_path = os.path.join(DATA_DIR, filename)
-    if os.path.exists(local_path):
-        return local_path
-    return f"{HF_BASE_URL}/{filename}"
-
-
 @app.on_event("startup")
 def startup():
     global db, registry
@@ -236,21 +224,17 @@ def startup():
 
     # Initialize DuckDB with resource limits suitable for CPU Basic.
     db = duckdb.connect(":memory:")
-    db.execute("INSTALL httpfs; LOAD httpfs;")
     db.execute(f"SET memory_limit='{DUCKDB_MEMORY_LIMIT}';")
     db.execute(f"SET threads={DUCKDB_THREADS};")
     db.execute(f"SET temp_directory='{DUCKDB_TEMP_DIR}';")
     db.execute("SET preserve_insertion_order=false;")
 
-    # Resolve data sources (remote or local)
-    results_src = _resolve_parquet("benchmark_results.parquet")
-    features_src = _resolve_parquet("feature_importance.parquet")
+    results_src = os.path.join(DATA_DIR, "benchmark_results.parquet")
+    features_src = os.path.join(DATA_DIR, "feature_importance.parquet")
 
-    is_remote = results_src.startswith("http")
-    mode = "remote httpfs" if is_remote else "local files"
     logger.info(
-        "DuckDB settings — memory_limit=%s threads=%d temp_dir=%s mode=%s",
-        DUCKDB_MEMORY_LIMIT, DUCKDB_THREADS, DUCKDB_TEMP_DIR, mode,
+        "DuckDB settings — memory_limit=%s threads=%d temp_dir=%s data_dir=%s",
+        DUCKDB_MEMORY_LIMIT, DUCKDB_THREADS, DUCKDB_TEMP_DIR, DATA_DIR,
     )
 
     # Load both tables into memory for reliable, fast queries.
@@ -277,22 +261,15 @@ def startup():
     n_features = db.execute("SELECT COUNT(*) FROM features").fetchone()[0]
     logger.info("Tables loaded: %d results, %d features", n_results, n_features)
 
-    # Load registry JSON (small file, always downloaded)
+    # Load registry JSON from local data dir.
     registry_local = os.path.join(DATA_DIR, "benchmark_registry.json")
-    if os.path.exists(registry_local):
+    try:
         with open(registry_local) as f:
             registry = json.load(f)
-    else:
-        # Download from HF
-        registry_url = f"{HF_BASE_URL}/benchmark_registry.json"
-        try:
-            with urllib.request.urlopen(registry_url) as resp:
-                registry = json.loads(resp.read().decode())
-            logger.info("Downloaded registry from %s", registry_url)
-        except Exception as e:
-            logger.warning("Could not load registry: %s", e)
-            registry = {"project": {}, "blood_platforms": {},
-                         "brain_targets": {}, "phases": {}, "models": {}}
+    except Exception as e:
+        logger.warning("Could not load registry from %s: %s", registry_local, e)
+        registry = {"project": {}, "blood_platforms": {},
+                     "brain_targets": {}, "phases": {}, "models": {}}
 
 
 @app.on_event("shutdown")