Add /health check based on RAM usage (#22)

Author: alejoe91

Dockerfile

@@ -5,6 +5,7 @@ WORKDIR /app
 ADD src ./src
 ADD pyproject.toml .
 ADD setup.py .
+ADD entrypoint.py .
 
 
 RUN apt-get update
@@ -31,10 +32,7 @@ RUN pip install scikit-learn==1.6.1
 # Install spikeinterface-gui from source
 RUN pip install spikeinterface-gui==0.13.1
 
-
-ENV PYTHONUNBUFFERED=1
-
 ENV PYTHONUNBUFFERED=1
 
 EXPOSE 8000
-ENTRYPOINT ["sh", "-c", "panel serve src/aind_ephys_portal/ephys_gui_app.py src/aind_ephys_portal/ephys_portal_app.py src/aind_ephys_portal/ephys_launcher_app.py src/aind_ephys_portal/ephys_monitor_app.py --setup src/aind_ephys_portal/setup.py --static-dirs images=src/aind_ephys_portal/images --address 0.0.0.0 --port 8000 --allow-websocket-origin ${ALLOW_WEBSOCKET_ORIGIN} --index ephys_portal_app.py --check-unused-sessions 2000 --unused-session-lifetime 5000 --num-threads 8"]
+ENTRYPOINT ["python", "entrypoint.py", "--address", "0.0.0.0", "--port", "8000", "--test"]

README.md

@@ -27,23 +27,15 @@ pip install -e .
 To run the Ephys Portal you can either use tha `launch.sh` script:
 
 ```bash
-bash launch.sh
-```
+AWS_PROFILE=your-profile
 
-or serve panel apps directly:
+# Make sure your SSO session is active
+aws sso login --profile $AWS_PROFILE
 
-```bash
-panel serve \
-    src/aind_ephys_portal/ephys_gui_app.py \
-    src/aind_ephys_portal/ephys_launcher_app.py \
-    src/aind_ephys_portal/ephys_monitor_app.py \
-    src/aind_ephys_portal/ephys_log_app.py \
-    --setup src/aind_ephys_portal/setup.py \
-    --static-dirs images=src/aind_ephys_portal/images \
-    --index ephys_portal_app.py \
-    --num-procs 4
+python entrypoint.py --port 8000 (default) --address 0.0.0.0 (default)
 ```
 
+
 This will start a Panel server and make the application available in your web browser.
 
 ## Development
@@ -58,7 +50,15 @@ pip install -e ".[dev]"
 1. Build the Docker image locally and run a Docker container:
 ```sh
 docker build -t aind-ephys-portal .
-docker run --rm -e ALLOW_WEBSOCKET_ORIGIN=0.0.0.0:8000 -v ~/.aws:/root/.aws:ro -p 8000:8000 aind-ephys-portal
+
+AWS_PROFILE=your-profile
+
+# Make sure your SSO session is active
+aws sso login --profile $AWS_PROFILE
+
+# Export temporary credentials and run Docker
+eval "$(aws configure export-credentials --profile $AWS_PROFILE --format env)"
+docker run --rm -e AWS_ACCESS_KEY_ID -e AWS_SECRET_ACCESS_KEY -e AWS_SESSION_TOKEN -e ALLOW_WEBSOCKET_ORIGIN=0.0.0.0:8000 -p 8000:8000 aind-ephys-portal
 ```
 2. Navigate to '0.0.0.0:8000` to view the app.
 

entrypoint.py

@@ -0,0 +1,102 @@
+import os
+import shutil
+from argparse import ArgumentParser
+
+import requests
+import psutil
+import panel as pn
+from tornado.web import RequestHandler
+
+# 1. Run setup (replaces --setup flag)
+from aind_ephys_portal.setup import *  # noqa: F401,F403
+from aind_ephys_portal.panel.logging import list_gui_sessions, get_max_number_of_gui_sessions, LOG_DIR  # noqa: F401
+
+
+if LOG_DIR.is_dir():
+    print(f"Cleaning up old log files in {LOG_DIR}...")
+    shutil.rmtree(LOG_DIR)
+
+
+def get_ecs_task_id():
+    metadata_uri = os.environ.get("ECS_CONTAINER_METADATA_URI_V4")
+    if metadata_uri:
+        # Request the metadata from the local ECS agent
+        response = requests.get(f"{metadata_uri}/task")
+        task_arn = response.json().get("TaskARN")
+        # The Task ID is the last part of the ARN string
+        return task_arn.split('/')[-1]
+    return "local-dev"
+
+
+MAX_GUI_SESSIONS_PER_TASK = get_max_number_of_gui_sessions()
+print(f"Max GUI sessions per task: {MAX_GUI_SESSIONS_PER_TASK}")
+
+# 2. Health Check & Index Redirect
+class HealthHandler(RequestHandler):
+    def get(self):
+        mem = psutil.virtual_memory()
+        # count number of GUI app sessions
+        gui_sessions = list_gui_sessions()
+        task_id = get_ecs_task_id()
+        if mem.percent > 70 or len(gui_sessions) > MAX_GUI_SESSIONS_PER_TASK:
+            self.set_status(503)
+            self.write(
+                f"Busy:\nMemory at {mem.percent}% - Num sessions: {len(gui_sessions)} "
+                f"(max {MAX_GUI_SESSIONS_PER_TASK}) Task ID: {task_id}"
+            )
+        else:
+            self.set_status(200)
+            self.write(
+                f"Healthy:\nMemory at {mem.percent}% - Num sessions: {len(gui_sessions)} "
+                f"(max {MAX_GUI_SESSIONS_PER_TASK}) Task ID: {task_id}"
+            )
+
+class IndexRedirectHandler(RequestHandler):
+    def get(self):
+        self.redirect("/ephys_portal_app")
+
+# 3. App file paths — Panel will exec these per-session with a proper context
+APP_DIR = "src/aind_ephys_portal"
+apps = {
+    "ephys_portal_app": os.path.join(APP_DIR, "ephys_portal_app.py"),
+    "ephys_gui_app": os.path.join(APP_DIR, "ephys_gui_app.py"),
+    "ephys_launcher_app": os.path.join(APP_DIR, "ephys_launcher_app.py"),
+    "ephys_monitor_app": os.path.join(APP_DIR, "ephys_monitor_app.py"),
+}
+
+# 4. Parse ALLOW_WEBSOCKET_ORIGIN from env
+allow_ws = os.environ.get("ALLOW_WEBSOCKET_ORIGIN", "*").split(",")
+
+parser = ArgumentParser(description="Ephys Portal Server")
+parser.add_argument("--port", type=int, default=8000, help="Port to run the server on")
+parser.add_argument("--address", type=str, default="localhost", help="Address to run the server on")
+parser.add_argument("--test", action="store_true", help="Run in test mode (connects to test API gateway)")
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    port = args.port
+    address = args.address
+    test_mode = args.test
+
+    # Set number of threads for Panel's thread pool to handle multiple sessions in parallel
+    pn.config.nthreads = 8
+
+    if test_mode:
+        print("Running in TEST MODE: connecting to test API gateway")
+        os.environ["TEST_ENV"] = "1"
+
+    print(f"Ephys Portal is running on http://{address}:{port}")
+    for app in apps:
+        print(f" - {app}: http://{address}:{port}/{app}")
+    print(f" - Health check: http://{address}:{port}/health")
+    pn.serve(
+        apps,
+        address=address,
+        port=port,
+        allow_websocket_origin=allow_ws,
+        static_dirs={"images": os.path.join(APP_DIR, "images")},
+        extra_patterns=[(r"/health", HealthHandler), (r"/", IndexRedirectHandler)],
+        check_unused_sessions=2000,
+        unused_session_lifetime=5000,
+        show=False,
+    )

launch.sh

@@ -6,26 +6,36 @@
 import panel as pn
 from aind_data_access_api.document_db import MetadataDbClient
 
+
+TEST_ENV = os.environ.get("TEST_ENV", "0") == "1"
+
 # Constants for database connection
-API_GATEWAY_HOST = os.environ.get("API_GATEWAY_HOST", "api.allenneuraldynamics-test.org")
+default_api_gateway_host = "api.allenneuraldynamics-test.org" if TEST_ENV else "api.allenneuraldynamics.org"
+API_GATEWAY_HOST = os.environ.get("API_GATEWAY_HOST", default_api_gateway_host)
 DATABASE = os.environ.get("DATABASE", "metadata_index")
 COLLECTION = os.environ.get("COLLECTION", "data_assets")
 
 # Timeouts
-TIMEOUT_1M = 60
 TIMEOUT_1H = 60 * 60
-TIMEOUT_24H = 60 * 60 * 24
 
 # Initialize the client
-client = MetadataDbClient(
+client_v1 = MetadataDbClient(
+    host=API_GATEWAY_HOST,
+    database=DATABASE,
+    collection=COLLECTION,
+    version="v1",
+)
+
+client_v2 = MetadataDbClient(
     host=API_GATEWAY_HOST,
     database=DATABASE,
     collection=COLLECTION,
+    version="v2",
 )
 
 
 @pn.cache()
-def get_name_from_id(id: str):
+def get_name_from_id(id: str, version: str = "v2") -> str:
     """Get the name field from a record with the given ID.
 
     Parameters
@@ -38,6 +48,7 @@ def get_name_from_id(id: str):
     str
         The name field from the record.
     """
+    client = client_v1 if version == "v1" else client_v2
     response = client.aggregate_docdb_records(pipeline=[{"$match": {"_id": id}}, {"$project": {"name": 1, "_id": 0}}])
     return response[0]["name"]
 
@@ -81,7 +92,7 @@ def get_asset_by_name(asset_name: str):
 
 
 @pn.cache(ttl=TIMEOUT_1H)
-def get_raw_asset_by_name(asset_name: str):
+def get_raw_asset_by_name(asset_name: str, version: str = "v2"):
     """Get all assets that match a given asset name pattern.
 
     Parameters
@@ -95,14 +106,15 @@ def get_raw_asset_by_name(asset_name: str):
         List of matching asset records.
     """
     raw_name = _raw_name_from_derived(asset_name)
+    client = client_v1 if version == "v1" else client_v2
     response = client.retrieve_docdb_records(
         filter_query={"name": {"$regex": raw_name, "$options": "i"}, "data_description.data_level": "raw"}, limit=0
     )
     return response
 
 
 @pn.cache(ttl=TIMEOUT_1H)
-def get_all_ecephys_derived(additional_includes_in_name: str | None = None) -> List[Dict[str, Any]]:
+def get_all_ecephys_derived(additional_includes_in_name: str | None = None, version: str = "v2") -> List[Dict[str, Any]]:
     """Get a limited set of all records from the database.
 
     Returns
@@ -113,6 +125,7 @@ def get_all_ecephys_derived(additional_includes_in_name: str | None = None) -> L
         Comma-separated list of additional fields to include in the results, by default None
     """
     filter_query = {"data_description.modality.abbreviation": "ecephys", "data_description.data_level": "derived"}
+    client = client_v1 if version == "v1" else client_v2
     responses = client.retrieve_docdb_records(
         filter_query=filter_query,
     )

src/aind_ephys_portal/docdb/database.py

@@ -17,7 +17,9 @@
 from spikeinterface.core.core_tools import extractor_dict_iterator, set_value_in_extractor_dict
 from spikeinterface.curation import validate_curation_dict
 
-from aind_ephys_portal.panel.logging import setup_logging, local_log_context
+from aind_ephys_portal.panel.logging import (
+    setup_logging, local_log_context, get_max_number_of_gui_sessions, list_gui_sessions
+)
 from aind_ephys_portal.panel.utils import PostMessageListener, FullscreenResizeHandler
 
 
@@ -94,7 +96,20 @@ def __init__(self, analyzer_path, recording_path, identifier=None, fast_mode=Fal
         self.loading_banner = pn.Row(self.spinner, self.log_output, sizing_mode="stretch_both")
 
         self.win = None
-        if self.analyzer_path != "":
+
+        num_gui_sessions = len(list_gui_sessions())
+        max_sessions = get_max_number_of_gui_sessions()
+        if num_gui_sessions > max_sessions:
+            print(f"Current number of GUI sessions: {num_gui_sessions}. Max allowed per worker: {max_sessions}.")
+            self.layout = pn.Column(
+                pn.pane.Markdown(
+                    f"⚠️ Too many active GUI sessions ({num_gui_sessions}). Max allowed per worker is {max_sessions}. "
+                    f"Please try again in a few minutes.",
+                    sizing_mode="stretch_both",
+                ),
+                sizing_mode="stretch_both",
+            )
+        elif self.analyzer_path != "":
             self.layout = pn.Column(
                 self._create_main_window(),
                 sizing_mode="stretch_both",