improved memory management and diagnostics (#36)

Author: alejoe91

Dockerfile

@@ -30,6 +30,10 @@ RUN pip install spikeinterface==0.104.1
 RUN pip install spikeinterface-gui==0.13.1
 
 ENV PYTHONUNBUFFERED=1
+# Limit glibc malloc arenas to reduce per-thread heap fragmentation.
+# Default is 8 * NCPU; with numpy/zarr large alloc + free patterns this
+# leaves freed pages stranded across many arenas, inflating RSS.
+ENV MALLOC_ARENA_MAX=2
 
 EXPOSE 8000
 ENTRYPOINT ["python", "entrypoint.py", "--address", "0.0.0.0", "--port", "8000"]

entrypoint.py

@@ -1,16 +1,25 @@
+import gc
 import os
 import shutil
 import threading
+import tracemalloc
 from argparse import ArgumentParser
+from collections import Counter
 
 import numpy as np
 import psutil
 import panel as pn
 from tornado.web import RequestHandler
 
+# Start tracemalloc as early as possible so allocation tracebacks include the
+# session-loading code paths. Frame depth 10 is enough to identify the call site
+# inside spikeinterface/zarr/fsspec without blowing up memory overhead.
+tracemalloc.start(10)
+_tracemalloc_baseline = tracemalloc.take_snapshot()
+
 # 1. Run setup (replaces --setup flag)
-from aind_ephys_portal.setup import *  # noqa: F401,F403
-from aind_ephys_portal.panel.logging import (
+from aind_ephys_portal.setup import *  # noqa: F401,F403,E402
+from aind_ephys_portal.panel.logging import (  # noqa: E402
     list_gui_sessions,
     get_max_number_of_gui_sessions,
     get_container_total_memory,
@@ -23,6 +32,13 @@
 TARGET_CLEAR_TMP_ARR_SECONDS = 180
 TARGET_INFLATE_DELAY_SECONDS = 90
 
+# Recycle signal: when a task is sitting idle (0 GUI sessions) but its RAM
+# stays above this percent, /health returns 503 so the ALB deregisters it and
+# the ECS service scheduler replaces it. Baseline (no sessions, fresh start)
+# is ~10%, so 40% means tolerating ~30 points of accumulated residue before
+# recycling. Only triggers with 0 sessions, so user sessions are never cut.
+RECYCLE_RAM_PERCENT_WHEN_IDLE = 40
+
 
 if LOG_DIR.is_dir():
     print(f"Cleaning up old log files in {LOG_DIR}...")
@@ -76,6 +92,16 @@ def get(self):
         task_id = get_ecs_task_id()
         max_gui_sessions = get_max_number_of_gui_sessions()
 
+        # 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.
+        if len(gui_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
+
         if mem_percent > TARGET_MEMORY_TRIGGER_PERCENT:
             self.set_status(200)
             self.write(
@@ -124,6 +150,99 @@ def get(self):
         self.redirect("/ephys_portal_app")
 
 
+class DebugMemoryHandler(RequestHandler):
+    """Live introspection for diagnosing memory retention on a bloated task.
+
+    Returns top object counts (via gc.get_objects), top allocation sites
+    (tracemalloc snapshot diff vs. boot baseline), and fsspec/s3fs instance
+    cache sizes. Gated by DEBUG_MEMORY_TOKEN env var to avoid exposing
+    process internals publicly.
+    """
+
+    def get(self):
+        token = os.environ.get("DEBUG_MEMORY_TOKEN")
+        if token and self.request.headers.get("X-Debug-Token") != token:
+            self.set_status(401)
+            self.write("Unauthorized: missing or wrong X-Debug-Token header")
+            return
+
+        task_id = get_ecs_task_id()
+        total_memory = get_container_total_memory()
+        used_memory = get_container_used_memory()
+        mem_percent = used_memory / total_memory * 100
+        gui_sessions = list_gui_sessions()
+
+        # Force a collection so we don't count obviously-dead objects.
+        gc.collect()
+        gc.collect()
+
+        # 1) Object counts by type (top 30 by count)
+        type_counts = Counter()
+        for obj in gc.get_objects():
+            type_counts[type(obj).__name__] += 1
+        top_types = type_counts.most_common(30)
+
+        # 2) numpy/zarr/pandas big buffers — sum bytes by type
+        big_buffers = []
+        try:
+            import numpy as _np
+
+            np_bytes = 0
+            np_count = 0
+            for obj in gc.get_objects():
+                if isinstance(obj, _np.ndarray):
+                    np_bytes += obj.nbytes
+                    np_count += 1
+            big_buffers.append(("numpy.ndarray", np_count, np_bytes))
+        except Exception as e:
+            big_buffers.append(("numpy.ndarray", -1, -1))
+
+        # 3) fsspec/s3fs cache state
+        fsspec_info = []
+        try:
+            from fsspec import AbstractFileSystem
+
+            fsspec_info.append(
+                f"fsspec AbstractFileSystem._cache size: {len(AbstractFileSystem._cache)}"
+            )
+            for key, fs in list(AbstractFileSystem._cache.items())[:10]:
+                fsspec_info.append(f"  - {type(fs).__name__}: protocol={getattr(fs, 'protocol', '?')}")
+        except Exception as e:
+            fsspec_info.append(f"fsspec inspection failed: {e}")
+
+        # 4) tracemalloc top allocators since boot baseline
+        tm_lines = []
+        try:
+            current = tracemalloc.take_snapshot()
+            stats = current.compare_to(_tracemalloc_baseline, "lineno")[:25]
+            for stat in stats:
+                tm_lines.append(
+                    f"  +{stat.size_diff/1024/1024:7.2f} MiB ({stat.count_diff:+d} blocks)  {stat.traceback[0]}"
+                )
+        except Exception as e:
+            tm_lines.append(f"tracemalloc unavailable: {e}")
+
+        self.set_header("Content-Type", "text/plain; charset=utf-8")
+        out = [
+            f"=== Task {task_id} ===",
+            f"Memory: {mem_percent:.1f}%  ({used_memory/1024**3:.2f} / {total_memory/1024**3:.2f} GB)",
+            f"GUI sessions: {len(gui_sessions)}",
+            "",
+            "--- Top object types by count ---",
+        ]
+        for name, count in top_types:
+            out.append(f"  {count:>10d}  {name}")
+        out.extend(["", "--- Buffer bytes ---"])
+        for name, count, nbytes in big_buffers:
+            mib = nbytes / 1024 / 1024 if nbytes >= 0 else -1
+            out.append(f"  {name}: {count} objects, {mib:.1f} MiB")
+        out.extend(["", "--- fsspec ---"])
+        out.extend(fsspec_info)
+        out.extend(["", "--- tracemalloc (top 25 since boot) ---"])
+        out.extend(tm_lines)
+        self.write("\n".join(out))
+
+
 # 3. App file paths — Panel will exec these per-session with a proper context
 APP_DIR = "src/aind_ephys_portal"
 apps = {
@@ -182,7 +301,11 @@ def get(self):
         port=port,
         allow_websocket_origin=allow_ws,
         static_dirs={"images": os.path.join(APP_DIR, "images")},
-        extra_patterns=[(r"/health", HealthHandler), (r"/", IndexRedirectHandler)],
+        extra_patterns=[
+            (r"/health", HealthHandler),
+            (r"/debug/memory", DebugMemoryHandler),
+            (r"/", IndexRedirectHandler),
+        ],
         check_unused_sessions=2000,
         unused_session_lifetime=5000,
         show=False,

src/aind_ephys_portal/panel/ephys_gui.py

@@ -23,10 +23,19 @@
     get_max_number_of_gui_sessions,
     list_gui_sessions,
     get_ecs_task_id,
+    get_container_total_memory,
+    get_container_used_memory,
     remove_session,
 )
 from aind_ephys_portal.panel.utils import PostMessageListener, FullscreenResizeHandler
 
+# Refuse to admit a new session if the task's RAM is already above this percent,
+# regardless of how many sessions are active. Protects "poisoned" tasks
+# (residue from prior sessions) from being pushed into OOM by the count-based
+# admission rule. Picked below the /health 503 recycle threshold so the ALB
+# pulls the task out of rotation before the GUI starts rejecting.
+MAX_RAM_PERCENT_FOR_NEW_SESSION = 55
+
 displayed_unit_properties = [
     "decoder_label",
     "default_qc",
@@ -78,6 +87,47 @@ def _malloc_trim():
         pass
 
 
+def _clear_fsspec_instance_caches():
+    """Drop cached fsspec filesystem instances and their internal block buffers.
+
+    fsspec keeps every filesystem ever constructed in AbstractFileSystem._cache,
+    so per-instance invalidate_cache() (dir listings) doesn't release the FS
+    itself or its dircache/block buffers. We only drop instances with no other
+    referrers to avoid stealing FS objects from concurrent sessions.
+    """
+    try:
+        import sys
+        from fsspec import AbstractFileSystem
+    except Exception:
+        return
+
+    cache = getattr(AbstractFileSystem, "_cache", None)
+    if not cache:
+        return
+
+    # 2 = the local var here + sys.getrefcount's own arg ref → no external holders.
+    # If anything else has a handle, leave it alone.
+    removed = 0
+    for key in list(cache.keys()):
+        fs = cache.get(key)
+        if fs is None:
+            continue
+        if sys.getrefcount(fs) <= 3:  # cache dict + local + getrefcount arg
+            try:
+                # Drop any caches the FS exposes before dropping the FS.
+                for attr in ("dircache", "_intrans", "_open_files"):
+                    try:
+                        getattr(fs, attr, {}).clear()
+                    except Exception:
+                        pass
+                del cache[key]
+                removed += 1
+            except Exception:
+                pass
+    if removed:
+        print(f"Dropped {removed} idle fsspec filesystem instance(s) from cache.")
+
+
 class EphysGuiView(param.Parameterized):
 
     def __init__(
@@ -126,20 +176,26 @@ def __init__(
 
         num_gui_sessions = len(list_gui_sessions())
         max_sessions = get_max_number_of_gui_sessions()
-        if num_gui_sessions > max_sessions:
+        ram_percent = get_container_used_memory() / get_container_total_memory() * 100
+        too_many_sessions = num_gui_sessions > max_sessions
+        too_much_ram = ram_percent > MAX_RAM_PERCENT_FOR_NEW_SESSION
+        if too_many_sessions or too_much_ram:
             # Remove this session from the count — it is being rejected
             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)
-            print(
-                f"Current number of GUI sessions: {num_gui_sessions}. Max allowed per worker: {max_sessions}. Task ID: {task_id}"
+            reason = (
+                f"too many sessions ({num_gui_sessions}/{max_sessions})"
+                if too_many_sessions
+                else f"task RAM already at {ram_percent:.1f}% (limit {MAX_RAM_PERCENT_FOR_NEW_SESSION}%)"
             )
+            print(f"Rejecting new session: {reason}. Task ID: {task_id}")
             self.layout = pn.Column(
                 header,
                 pn.pane.Markdown(
-                    f"⚠️ Too many active GUI sessions ({num_gui_sessions}). Max allowed per worker is {max_sessions}. "
+                    f"⚠️ Cannot start a new GUI session: {reason}. "
                     f"Please try again in a few minutes or open a new tab (Task ID: `{task_id}`).",
                     sizing_mode="stretch_both",
                 ),
@@ -543,6 +599,8 @@ def cleanup(self):
             def _deferred_gc():
                 gc.collect()
                 gc.collect()
+                _clear_fsspec_instance_caches()
+                gc.collect()
                 _malloc_trim()
                 final_mem = psutil.virtual_memory()
                 used = final_mem.used / (1024**3)
@@ -554,6 +612,8 @@ def _deferred_gc():
             # Fallback: run immediately if IOLoop is unavailable
             gc.collect()
             gc.collect()
+            _clear_fsspec_instance_caches()
+            gc.collect()
             _malloc_trim()
             final_mem = psutil.virtual_memory()
             used = final_mem.used / (1024**3)