fix(storage#741): drop page cache before read_threads memory guard

[FileSystemGuy]

Problem

ConfigArguments.validate in utils/config.py samples
psutil.virtual_memory().available and caps read_threads against it
(the mlcommons/storage#448 per-node budget fix). On POSIX filesystems
whose client pins a large reclaimable page cache (Lustre in
particular), MemAvailable understates by the size of that
cache — so after run 1 warms the cache, runs 2..N of a 5-run
submission see the cache as "unavailable" and DLIO refuses to launch:

Per-node memory budget exceeded on host X: reader.read_threads=17 x
local_ranks=16 = 288 worker processes, estimated ~144 GB (available
RAM on this node: 159 GB; total: 188 GB). Reduce reader.read_threads
to at most 16.

This aborts initialize() before the per-epoch page-cache flush
ever gets to run. reportgen then reports INVALID: 0 runs for a
submission whose only real problem is that the benchmark's own
per-epoch flush (which drops that exact cache every epoch) hadn't run
yet. Reporter's evidence: run 1 passes at ~170 GB MemAvailable; runs
2–5 crash at ~147–150 GB; sync; echo 1 > drop_caches restores ~174
GB and the guard passes again — confirming the "used" RAM is
reclaimable cache.

Filed as mlcommons/storage#741 (@mirajeev).

Fix

Right before sampling virtual_memory(), drop the local host's page
cache using the same sudo -n sh -c 'echo 3 > /proc/sys/vm/drop_caches'
command already invoked per-epoch in main.py. Gated on
local_rank() == 0 so exactly one flush runs per host, followed by a
comm().Barrier() so non-leader ranks read post-drop memory.

Design choices worth flagging on review

local_rank() not MPI.node(). MPI.node() is computed from
cumulative offsets of mpi_ppn_list (utility.py:313–318) and assumes
contiguous per-node rank blocks. Under --map-by node (which
mlpstorage sets by default for multi-host) that assumption is
violated: rank 0 → host 0, rank 1 → host 1, …, rank N → host 0, and
two different physical hosts end up computing the same node index.
This is the same anti-pattern PR mlcommons/storage#675 fixed for the
host_memory_GB collector array (the [376.18 × 7, 187.90, 0.0 × 7]
shape). local_rank() derives from MPI.COMM_TYPE_SHARED — the
MPI-3 primitive that groups ranks by shared memory — and is
independent of rank assignment. Worst-case failure mode is
redundant/idempotent flushes on a host, never a missed flush on a host
that needed one.

No SSH dependency. The flush runs inside the DLIO process, which
is already up on each host via whatever launcher was used (OpenMPI,
PALS palsd, Slurm slurmstepd). We piggyback on the transport that
already worked. In particular this is compatible with the
--skip-ssh-check / PALS/Slurm auto-skip path in
mlcommons/storage#740 — those users need this fix precisely because
their sites disallow passwordless SSH between compute nodes, so any
mlpstorage-side fan-out via mpirun would have been unusable for
them.

Fail-open matches per-epoch flush. Every failure mode — sudo -n
refused, kernel timeout, Barrier() error, subprocess OSError —
is swallowed. The guard downstream still runs; the worst case is that
it still trips on this host, which is exactly current behavior. No
regression path.

Same env-var timeout knob. Uses _resolve_drop_caches_timeout()
(30s default, DLIO_DROP_CACHES_TIMEOUT override). Operators already
tune this for the per-epoch flush; introducing a second knob would
double the operational surface for no gain.

Skipped when MPI not initialized. In child processes and test
harnesses mpi_state != MPI_INITIALIZED; the flush is skipped so
comm() isn't called and the guard still runs unchanged.

Test plan

• 18 new unit tests in tests/test_drop_caches_before_memory_guard.py:
  MPI-state gating, local_rank() gate (leader + non-leader
  parametrized), fail-open across every subprocess exception class
  + non-zero exit + Barrier failure on leader/non-leader, timeout
  plumbing (default + env override + bad-env-value).
• Existing tests/test_drop_caches_timeout.py (22 tests) still
  passes — no changes to _resolve_drop_caches_timeout.
• uv run pytest tests/ --ignore=<benchmark end-to-end suites> —
  183 passed, 2 skipped, no regressions.
• uv run pytest tests/test_fast_ci.py — 92 passed, 1 skipped.
  Includes TestEndToEndSmoke::test_train_npy_smoke which
  exercises the modified validate() path with read_threads > 0
  and PyTorch data loader.

Related

• Reporter's evidence + full reproduction: RAM budget guard fails consecutive training runs — sizes read_threads from MemAvailable, which under-counts reclaimable page cache storage#741
• Filed by @mirajeev alongside a between-runs cache-drop workaround
• Related to Worker-memory budget check compares per-node RAM against global (comm_size-wide) worker count → false positives at scale storage#448 (per-node vs world-comm scoping fix
  this guard was originally built for)
• SSH-less-launcher context: feat(validation): --skip-ssh-check + auto-skip SSH preflight for PALS/Slurm launchers storage#740
• MPI.node() unreliability motivating local_rank() gate:
  Bug: Submission checker rule 3.1.2 double-counts host memory, doubling the required dataset size storage#669 / PR fix(#669): aggregate host_memory_GB via sum(), not num_hosts × array[0] storage#675
• Shared timeout mechanism: Could not flush page cache between epochs warning storage#487 / DLIO PR fix(main): handle slow flush as transient; expose DLIO_DROP_CACHES_TIMEOUT #28