Dev changes to main (#47)

Author: alejoe91

Dockerfile

@@ -23,11 +23,14 @@ RUN wget "https://www.wavpack.com/wavpack-${WAVPACK_VERSION}.tar.bz2" && \
 # Install
 RUN pip install wavpack-numcodecs
 
-# Install spikeinterface from source
+# Install spikeinterface
 RUN pip install spikeinterface==0.104.1
 
-# Install spikeinterface-gui from source
-RUN pip install spikeinterface-gui==0.13.1
+# Install spikeinterface-gui form branch with AIND fixes
+RUN git clone https://github.com/alejoe91/spikeinterface-gui.git && \
+    cd spikeinterface-gui && git checkout e98d515bb25c10a357c7eee3afb85e55d7afab67 && \
+    pip install . && cd ..
+# RUN pip install spikeinterface-gui==0.13.1
 
 # Pin scikit-learn AFTER spikeinterface installs so we override whatever the
 # transitive resolver picked. Match the version that analyzers in our pipeline
@@ -43,4 +46,4 @@ ENV PYTHONUNBUFFERED=1
 ENV MALLOC_ARENA_MAX=2
 
 EXPOSE 8000
-ENTRYPOINT ["python", "entrypoint.py", "--address", "0.0.0.0", "--port", "8000"]
+ENTRYPOINT ["python", "entrypoint.py", "--address", "0.0.0.0", "--port", "8000"]

entrypoint.py

@@ -30,16 +30,18 @@
 
 # 1. Run setup (replaces --setup flag)
 from aind_ephys_portal.setup import *  # noqa: F401,F403,E402
-from aind_ephys_portal.panel.logging import (  # noqa: E402
+from aind_ephys_portal.session_logging import (  # noqa: E402
     list_gui_sessions,
     get_max_number_of_gui_sessions,
-    can_admit_new_session,
     get_hard_cap_sessions,
+    get_health_estimate_pct,
     get_container_total_memory,
     get_container_used_memory,
     get_ecs_task_id,
+    recent_session_rejection,
     LOG_DIR,
 )  # noqa: F401
+from aind_ephys_portal.ecs_protection import is_protected, get_protection_status  # noqa: E402
 
 TARGET_MEMORY_TRIGGER_PERCENT = 70
 TARGET_CLEAR_TMP_ARR_SECONDS = 180
@@ -111,19 +113,25 @@ def get(self):
         # Idle-but-bloated → ask ALB to deregister so ECS replaces us.
         # GUI-level enforcement protects against admitting too many sessions,
         # but only the load balancer can take a leaky task out of rotation.
+        protected = is_protected()
+        prot_tag = " (protected)" if protected else ""
+
         if num_sessions == 0 and mem_percent > RECYCLE_RAM_PERCENT_WHEN_IDLE:
             self.set_status(503)
             self.write(
                 f"Unhealthy (recycle): Memory at {mem_percent:.1f}% with 0 sessions - Task ID: {task_id}"
             )
             return
 
-        # "Operationally full" = at least one session AND can't fit another
-        # without crossing the safe RAM ceiling or hitting the hard count cap.
-        # This unifies the old (count-based) and RAM-based busy paths into one
-        # predicate that's accurate for both light and heavy sessions.
-        full = num_sessions > 0 and not can_admit_new_session(
-            current_count=num_sessions, used_pct=mem_percent
+        # "Operationally full" = either at the hard session cap, or the GUI
+        # just rejected an incoming session (RAM headroom exceeded). We do
+        # NOT predict-reject here based on a hypothetical incoming session —
+        # that would over-trigger ECS scale-up whenever a mid/heavy session
+        # is happily running. Concrete rejections are the only signal that
+        # real scaling pressure exists.
+        full = num_sessions > 0 and (
+            num_sessions >= get_hard_cap_sessions()
+            or recent_session_rejection()
         )
 
         if full:
@@ -150,7 +158,7 @@ def get(self):
             elif _tmp_array_triggered:
                 busy_msg += " (inflated memory released)"
             self.write(
-                f"Busy {busy_msg}:\nMemory at {mem_percent:.1f}% - Num sessions: {num_sessions} "
+                f"Busy {busy_msg}{prot_tag}:\nMemory at {mem_percent:.1f}% - Num sessions: {num_sessions} "
                 f"(hard cap {hard_cap}) Task ID: {task_id}"
             )
         else:
@@ -160,11 +168,22 @@ def get(self):
                 _inflate_delay_timer = None
             self.set_status(200)
             self.write(
-                f"Healthy:\nMemory at {mem_percent:.1f}% - Num sessions: {num_sessions} "
+                f"Healthy{prot_tag}:\nMemory at {mem_percent:.1f}% - Num sessions: {num_sessions} "
                 f"(hard cap {hard_cap}) Task ID: {task_id}"
             )
 
 
+class ProtectionStatusHandler(RequestHandler):
+    """GET /debug/protection — show current ECS task scale-in protection status."""
+
+    def get(self):
+        import json as _json
+
+        status = get_protection_status()
+        self.set_header("Content-Type", "application/json; charset=utf-8")
+        self.write(_json.dumps(status, indent=2))
+
+
 class IndexRedirectHandler(RequestHandler):
     def get(self):
         self.redirect("/ephys_portal_app")
@@ -425,6 +444,8 @@ def post(self):
     if test_mode:
         print("Running in TEST MODE: connecting to test API gateway")
         os.environ["TEST_ENV"] = "1"
+        # Lower threshold on test for easier testing
+        os.environ["RECYCLE_RAM_PERCENT_WHEN_IDLE"] = "15"
 
     print(f"Ephys Portal is running on http://{address}:{port}")
     for app in apps:
@@ -441,6 +462,7 @@ def post(self):
             (r"/debug/memory", DebugMemoryHandler),
             (r"/debug/tracemalloc/start", TracemallocStartHandler),
             (r"/debug/tracemalloc/stop", TracemallocStopHandler),
+            (r"/debug/protection", ProtectionStatusHandler),
             (r"/", IndexRedirectHandler),
         ],
         check_unused_sessions=2000,

src/aind_ephys_portal/ecs_protection.py

@@ -0,0 +1,134 @@
+"""ECS task scale-in protection via the container agent endpoint.
+
+When at least one GUI session is active the task is marked as protected so
+that ECS auto-scaling and deployment scale-in events cannot terminate it.
+Protection is renewed periodically (before the expiry window closes) and
+cleared when the last GUI session ends.
+
+The module is a no-op when ``ECS_AGENT_URI`` is not set (local development).
+"""
+
+import json
+import os
+import threading
+
+import requests
+
+from aind_ephys_portal.session_logging import list_gui_sessions
+
+# --- Configuration -----------------------------------------------------------
+
+PROTECTION_DURATION_MINUTES = 60
+"""Each protection window lasts this long. Renewed before it expires."""
+
+_REFRESH_INTERVAL_SECONDS = 50 * 60  # 50 min — well before 60-min expiry
+
+# --- Module state (guarded by _lock) -----------------------------------------
+
+_lock = threading.Lock()
+_is_protected = False
+_refresh_timer: threading.Timer | None = None
+
+
+# --- Low-level ECS agent call ------------------------------------------------
+
+def _agent_uri() -> str | None:
+    return os.environ.get("ECS_AGENT_URI")
+
+
+def _set_task_protection(enabled: bool, expires_minutes: int = PROTECTION_DURATION_MINUTES) -> bool:
+    """PUT to the ECS agent endpoint. Returns True on success."""
+    uri = _agent_uri()
+    if not uri:
+        action = "protect" if enabled else "unprotect"
+        print(f"[ecs_protection] No ECS_AGENT_URI — skipping {action} (local dev)")
+        return True  # treat as success in local dev
+
+    url = f"{uri}/task-protection/v1/state"
+    body: dict = {"ProtectionEnabled": enabled}
+    if enabled:
+        body["ExpiresInMinutes"] = expires_minutes
+
+    try:
+        resp = requests.put(url, json=body, timeout=5)
+        data = resp.json()
+        if "error" in data:
+            print(f"[ecs_protection] Agent error: {json.dumps(data['error'])}")
+            return False
+        print(f"[ecs_protection] Task protection set: enabled={enabled}, response={json.dumps(data)}")
+        return True
+    except Exception as exc:
+        print(f"[ecs_protection] Failed to set protection (enabled={enabled}): {exc}")
+        return False
+
+
+# --- High-level API ----------------------------------------------------------
+
+def _cancel_refresh_timer():
+    global _refresh_timer
+    if _refresh_timer is not None:
+        _refresh_timer.cancel()
+        _refresh_timer = None
+
+
+def _schedule_refresh():
+    global _refresh_timer
+    _cancel_refresh_timer()
+    _refresh_timer = threading.Timer(_REFRESH_INTERVAL_SECONDS, _refresh_protection)
+    _refresh_timer.daemon = True
+    _refresh_timer.start()
+
+
+def _refresh_protection():
+    """Called by the timer — renew protection if sessions still exist, else clear."""
+    with _lock:
+        if len(list_gui_sessions()) > 0:
+            print("[ecs_protection] Renewing task protection (sessions still active)")
+            _set_task_protection(True)
+            _schedule_refresh()
+        else:
+            print("[ecs_protection] No sessions at refresh time — clearing protection")
+            _unprotect_task_locked()
+
+
+def _unprotect_task_locked():
+    """Must be called while holding _lock."""
+    global _is_protected
+    _cancel_refresh_timer()
+    _set_task_protection(False)
+    _is_protected = False
+
+
+def protect_task():
+    """Enable scale-in protection for this task (idempotent while protected)."""
+    global _is_protected
+    with _lock:
+        if _is_protected:
+            return
+        if _set_task_protection(True):
+            _is_protected = True
+            _schedule_refresh()
+
+
+def unprotect_task():
+    """Disable scale-in protection for this task."""
+    with _lock:
+        if not _is_protected:
+            return
+        _unprotect_task_locked()
+
+
+def is_protected() -> bool:
+    return _is_protected
+
+
+def get_protection_status() -> dict | None:
+    """GET the current protection status from the ECS agent. Returns None in local dev."""
+    uri = _agent_uri()
+    if not uri:
+        return {"local_dev": True, "is_protected": _is_protected}
+    try:
+        resp = requests.get(f"{uri}/task-protection/v1/state", timeout=5)
+        return resp.json()
+    except Exception as exc:
+        return {"error": str(exc)}

src/aind_ephys_portal/ephys_launcher_app.py

@@ -4,7 +4,7 @@
 
 import panel as pn
 
-from aind_ephys_portal.panel.logging import setup_logging
+from aind_ephys_portal.session_logging import setup_logging
 from aind_ephys_portal.panel.utils import EPHYSGUI_LINK_PREFIX
 
 pn.extension()

src/aind_ephys_portal/ephys_monitor_app.py

@@ -2,7 +2,7 @@
 import pandas as pd
 import panel as pn
 
-from aind_ephys_portal.panel.logging import (
+from aind_ephys_portal.session_logging import (
     list_sessions,
     remove_session,
     get_container_total_memory,

src/aind_ephys_portal/panel/ephys_gui.py

@@ -35,7 +35,7 @@
 from spikeinterface.core.zarrextractors import super_zarr_open
 from spikeinterface.curation import validate_curation_dict
 
-from aind_ephys_portal.panel.logging import (
+from aind_ephys_portal.session_logging import (
     setup_logging,
     local_log_context,
     can_admit_new_session,
@@ -47,6 +47,11 @@
     get_ecs_task_id,
     get_container_total_memory,
     get_container_used_memory,
+    record_session_rejection,
+    record_pending_load,
+    clear_pending_load,
+    record_admission,
+    admission_cooldown_remaining,
     remove_session,
 )
 from aind_ephys_portal.panel.utils import PostMessageListener, FullscreenResizeHandler
@@ -211,6 +216,34 @@ def __init__(
         ram_percent = get_container_used_memory() / get_container_total_memory() * 100
         total_ram_bytes = get_container_total_memory()
 
+        # Admission cooldown: if another session was admitted within the
+        # last few seconds, soft-reject this one. The previous admission's
+        # pending-load entry needs a moment to register so concurrent
+        # admissions don't all read the same stale `used + pending` sum.
+        # NB: this is sequencing, not real capacity pressure — we do NOT
+        # call record_session_rejection() so /health doesn't inflate.
+        cooldown_remaining = admission_cooldown_remaining()
+        if cooldown_remaining > 0:
+            print(
+                f"Soft-rejecting (admission cooldown): {cooldown_remaining:.1f}s remaining. "
+                f"Task ID: {task_id}"
+            )
+            doc = pn.state.curdoc
+            route = getattr(doc, "_log_route", None)
+            session_id = getattr(doc, "_log_session_id", None)
+            if route and session_id:
+                remove_session(route, session_id)
+            self.layout = pn.Column(
+                header,
+                pn.pane.Markdown(
+                    f"⏱️ Another session is being admitted — please retry in "
+                    f"a few seconds. (Task: `{task_id}`)",
+                    sizing_mode="stretch_both",
+                ),
+                sizing_mode="stretch_both",
+            )
+            return
+
         # Try to derive a *dataset-specific* RAM estimate by reading the unit counts
         # unit count. This costs one S3 round-trip (~1-3s) but lets us accurately predict heavy
         # sessions instead of relying on a fixed percent.
@@ -227,6 +260,11 @@ def __init__(
                     f"{est_bytes / (1024**3):.2f} GB ({estimate_pct:.1f}%) "
                     f"for {num_units} units, fast_mode={self.fast_mode}"
                 )
+                # Teach /health what the largest recent session looked like
+                # — its admission predicate uses this so it can fire the
+                # inflate signal when a *heavy* session arrives instead of
+                # relying on the static fallback percent.
+                record_session_estimate(estimate_pct)
             except Exception as e:
                 # Pre-load failed (S3 transient, bad path, etc.) — fall back
                 # to the static estimate. Better to occasionally reject a
@@ -240,6 +278,11 @@ def __init__(
             used_pct=ram_percent,
             estimate_pct=estimate_pct,
         ):
+            # Tell /health a rejection just happened so it can fire the
+            # inflate / scale-up signal on its next tick. Without this,
+            # /health would see the task's current RAM and think there's
+            # still room — even though admission just turned a session away.
+            record_session_rejection()
             # Remove this session from the count — it is being rejected
             doc = pn.state.curdoc
             route = getattr(doc, "_log_route", None)
@@ -269,6 +312,19 @@ def __init__(
                 sizing_mode="stretch_both",
             )
         elif self.analyzer_path != "":
+            # Admission succeeded — register this session's estimated RAM so
+            # later admission attempts (within the same ~10s loading window)
+            # see it as "pending" instead of reading a stale-low used_pct from
+            # cgroup. Falls back to the static estimate if the dynamic one
+            # could not be computed. The pending entry is cleared in cleanup()
+            # or expires after PENDING_LOAD_TTL (60s).
+            self._pending_session_id = getattr(pn.state.curdoc, "_log_session_id", None)
+            pending_pct = estimate_pct if estimate_pct is not None else get_per_session_estimate_pct()
+            record_pending_load(self._pending_session_id, pending_pct)
+            # Start the admission-cooldown clock so the next arrival waits
+            # for THIS pending-load entry to be visible.
+            record_admission()
+
             self.layout = pn.Column(
                 header,
                 self._create_main_window(),
@@ -517,6 +573,13 @@ def cleanup(self):
         current_ram_usage = initial_mem.used / (1024**3)
         print(f"\nRAM Usage before cleanup: {current_ram_usage:.2f} / {total_ram:.2f} GB\n")
 
+        # 0) Release this session's pending-load reservation so the next
+        # admission attempt sees accurate RAM headroom. (The actual cgroup
+        # `memory.current` won't drop until the deferred GC further down
+        # runs malloc_trim — but the pending entry was overestimating, so
+        # dropping it now lets other admissions proceed sooner.)
+        clear_pending_load(getattr(self, "_pending_session_id", None))
+
         # 1) Clear postMessage listeners and submit trigger (they hold bound-method back-refs to self)
         self.curation_listener = None
         self.fullscreen_listener = None