fix: cap the memmap LRU below the fd limit, and let callers subset cases

Two failures found by running the loader over the real archive.

The memmap LRU holds an open file descriptor per entry and defaults to
1000, which sits just under the usual 1024 soft RLIMIT_NOFILE. Past that,
caching fails with a warning rather than an exception and every read falls
back to recomputing from the source, so a cached dataset silently becomes
an uncached one -- 50,176 such warnings in one run here. The size is now
capped at 75 % of the actual limit, with a warning naming ulimit -n.

VitalDBDataset also had no way to restrict membership: `target` decides
labels only, so a caller holding a 40-case cohort or a train split still
loaded all 5,269 cases. Adds a `cases` whitelist -- needed for training on
a split at all, not just for probes.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: Guido Gagliardi

physioex/data/base.py

@@ -162,7 +162,7 @@ def __init__(
         # which are expensive on NFS (~0.16ms per call on network filesystems).
         # Cache key is (subject_id, physical_channel, pipeline_hash).
         self._memmap_cache: OrderedDict[str, np.ndarray] = OrderedDict()
-        self._memmap_cache_size = int(memmap_cache_size)
+        self._memmap_cache_size = self._safe_memmap_cache_size(memmap_cache_size)
         self._memmap_cache_hits = 0
         self._memmap_cache_misses = 0
 
@@ -203,6 +203,40 @@ def __init__(
     # Subclass hooks
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _safe_memmap_cache_size(requested: int) -> int:
+        """Cap the memmap LRU so it cannot exhaust the process file limit.
+
+        Every cached memmap holds an open file descriptor.  The default of 1000
+        sits just under the usual 1024 soft ``RLIMIT_NOFILE``, so a dataset with
+        more than ~1000 (subject, channel) pairs runs the process out of
+        descriptors.  The symptom is not an exception: caching fails with a
+        warning and every read falls back to recomputing from the source file,
+        turning a cached dataset into an uncached one at a fraction of the speed,
+        silently.
+
+        Leaves at least a quarter of the budget for everything else the process
+        needs to open -- source files, the cache writes themselves, sockets.
+        """
+        try:
+            import resource
+
+            soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+        except Exception:  # pragma: no cover - non-POSIX
+            return int(requested)
+        if soft in (resource.RLIM_INFINITY, -1):
+            return int(requested)
+
+        budget = max(16, int(soft * 0.75))
+        if requested > budget:
+            logger.warning(
+                f"memmap_cache_size={requested} exceeds the safe budget for a "
+                f"file-descriptor limit of {soft}; capping at {budget}. Raise "
+                f"the limit (ulimit -n) to cache more."
+            )
+            return budget
+        return int(requested)
+
     @abstractmethod
     def _list_subjects(self) -> List[SubjectSpec]:
         """Return the subject manifest.  Lightweight: must NOT load signal data."""

physioex/data/datasets/vitaldb.py

@@ -39,7 +39,7 @@
 import csv
 import logging
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 
 import numpy as np
 
@@ -68,6 +68,10 @@ class VitalDBDataset(BasePhysioDataset):
             (the sentinel ignored by ``CrossEntropyLoss(ignore_index=-1)``).
         ane_type: keep only cases with this ``ane_type`` in ``cases.csv``.
             Defaults to ``"General"``; pass ``None`` to keep every case.
+        cases: restrict the dataset to these case ids.  Without it every case
+            on disk is listed, so a caller holding a cohort or a split would
+            silently load the whole archive; ``target`` only decides labels,
+            not membership.
         require_tracks: cases must carry all of these tracks (checked against
             ``meta/trks.csv``, no file opening).  Defaults to the EEG channel.
         epoch_length_sec: defaults to 120 s -- the 2-minute window used by the
@@ -106,6 +110,7 @@ def __init__(
         root: Optional[str] = None,
         target: TargetSpec = None,
         ane_type: Optional[str] = "General",
+        cases: Optional[Iterable[str]] = None,
         require_tracks: tuple = (EEG1,),
         **kwargs,
     ):
@@ -116,6 +121,7 @@ def __init__(
         # _list_subjects() and then _read_subject_labels().
         self._target = target
         self._ane_type = ane_type
+        self._case_filter = {str(c) for c in cases} if cases is not None else None
         self._require_tracks = tuple(require_tracks)
         self._cases: Dict[str, Dict[str, Any]] = {}
         # caseid -> resolved target, filled eagerly in _list_subjects
@@ -186,6 +192,8 @@ def _list_subjects(self) -> List[SubjectSpec]:
 
         subjects: List[SubjectSpec] = []
         for caseid, row in self._cases.items():
+            if self._case_filter is not None and caseid not in self._case_filter:
+                continue
             if self._ane_type is not None and row.get("ane_type") != self._ane_type:
                 continue
             if with_tracks is not None and caseid not in with_tracks:

tests/test_vitaldb_dataset.py

@@ -298,3 +298,40 @@ def test_real_vitaldb_smoke():
     assert ds.get_n_subjects() > 5_000
     item = ds[0]
     assert torch.is_tensor(item["labels"])
+
+
+def test_cases_whitelist_restricts_membership(vitaldb_root, tmp_path):
+    """Without this, holding a cohort or a split still loads every case.
+
+    ``target`` only decides labels: a caller that passed a 40-case cohort and
+    expected a 40-case dataset would quietly get the whole archive.
+    """
+    ds = make_dataset(vitaldb_root, tmp_path / "c1", cases=["1", "3"])
+    assert ds.get_subjects() == ["1", "3"]
+    assert make_dataset(vitaldb_root, tmp_path / "c2").get_subjects() == ["1", "2", "3"]
+
+
+def test_cases_whitelist_ignores_ids_not_on_disk(vitaldb_root, tmp_path):
+    ds = make_dataset(vitaldb_root, tmp_path / "cache", cases=["1", "9999"])
+    assert ds.get_subjects() == ["1"]
+
+
+def test_memmap_cache_is_capped_below_the_file_descriptor_limit(vitaldb_root, tmp_path):
+    """Every cached memmap holds an fd; the default 1000 sits under a 1024 limit.
+
+    Exhausting them does not raise -- caching fails with a warning and every
+    read recomputes from source, so a cached dataset silently becomes an
+    uncached one.
+    """
+    import resource
+
+    soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
+    ds = make_dataset(vitaldb_root, tmp_path / "cache", memmap_cache_size=1_000_000)
+    if soft not in (resource.RLIM_INFINITY, -1):
+        assert ds._memmap_cache_size < soft
+        assert ds._memmap_cache_size == max(16, int(soft * 0.75))
+
+
+def test_a_modest_memmap_cache_size_is_left_alone(vitaldb_root, tmp_path):
+    ds = make_dataset(vitaldb_root, tmp_path / "cache", memmap_cache_size=8)
+    assert ds._memmap_cache_size == 8