diff --git a/src/flashpack/parallel_read.py b/src/flashpack/parallel_read.py index 16401c0..3886063 100644 --- a/src/flashpack/parallel_read.py +++ b/src/flashpack/parallel_read.py @@ -41,6 +41,8 @@ - ``FLASHPACK_READ_CHUNK_BYTES`` chunk size (default 64 MiB) - ``FLASHPACK_DIRECT_IO=0`` never use O_DIRECT - ``FLASHPACK_CACHE_PINNED=0`` free pinned staging buffers after load (CUDA) +- ``FLASHPACK_SAMPLE_PROBE=0`` trust mincore alone for the O_DIRECT gate +- ``FLASHPACK_SAMPLE_PROBE_MIN_GBPS`` hot threshold for the sample probe (default 5.0) """ import ctypes @@ -48,6 +50,7 @@ import os import queue import threading +import time from typing import TYPE_CHECKING import numpy as np @@ -185,6 +188,81 @@ def _page_cache_resident_fraction(path: str, size: int) -> float: return 0.0 +_SAMPLE_PROBE_THREADS = 4 +_SAMPLE_PROBE_BYTES = 16 * 1024 * 1024 + + +def _sample_read_gbps(path: str, size: int) -> float: + """Aggregate buffered read rate over samples spread across the file + (GB/s; 0.0 on any failure). + + Discriminates page-hot from cache-cold where mincore cannot: hot files + scale with reader threads (well above 5 GB/s aggregate), while cold + network/NVMe-cache buffered reads saturate around 2-4 GB/s regardless + of thread count. Cost: at most 64 MiB of buffered reads, which land in + the page cache and are not wasted. + """ + if size <= _SAMPLE_PROBE_THREADS * _SAMPLE_PROBE_BYTES: + # Small file: buffered is the right choice either way. + return float("inf") + try: + results = [0] * _SAMPLE_PROBE_THREADS + step = (size - _SAMPLE_PROBE_BYTES) // (_SAMPLE_PROBE_THREADS - 1) + + def _sampler(idx: int) -> None: + fd = os.open(path, os.O_RDONLY) + try: + off = (idx * step) & ~(_ALIGN - 1) + want = min(_SAMPLE_PROBE_BYTES, size - off) + buf = bytearray(want) + got = os.preadv(fd, [buf], off) + results[idx] = max(got, 0) + finally: + os.close(fd) + + threads = [ + threading.Thread(target=_sampler, args=(i,), daemon=True) + for i in range(_SAMPLE_PROBE_THREADS) + ] + t0 = time.perf_counter() + for t in threads: + t.start() + for t in threads: + t.join() + dt = time.perf_counter() - t0 + total = sum(results) + if dt <= 0 or total == 0: + return 0.0 + return total / 1e9 / dt + except Exception: + return 0.0 + + +def _sample_probe_min_gbps() -> float: + try: + return float(os.environ.get("FLASHPACK_SAMPLE_PROBE_MIN_GBPS", "5.0")) + except ValueError: + return 5.0 + + +def _should_use_direct(path: str, size: int) -> bool: + """Decide O_DIRECT vs buffered for this load. + + mincore is the fast positive signal, but under cgroup-managed runners it + can report 0.0 residency for a demonstrably hot page cache (measured: + buffered repeat rode the cache at 21.6 GB/s right after mincore read 0.0), + which silently forces every warm reload onto the ~2x-slower direct path. + A timed sample read verifies coldness before O_DIRECT is chosen. + """ + if not (_env_flag("FLASHPACK_DIRECT_IO") and hasattr(os, "O_DIRECT")): + return False + if _page_cache_resident_fraction(path, size) >= 0.9: + return False + if not _env_flag("FLASHPACK_SAMPLE_PROBE", default=True): + return True + return _sample_read_gbps(path, size) < _sample_probe_min_gbps() + + def _read_chunk(fd_direct, fd_plain: int, view, f_off: int, ln: int) -> None: """Fill ``view[:ln]`` from file offset ``f_off``. @@ -267,11 +345,7 @@ def _parallel_read_into_cpu_storage( work.put(None) size = os.path.getsize(path) - use_direct = ( - _env_flag("FLASHPACK_DIRECT_IO") - and hasattr(os, "O_DIRECT") - and _page_cache_resident_fraction(path, size) < 0.9 - ) + use_direct = _should_use_direct(path, size) errors: list[BaseException] = [] @@ -350,13 +424,9 @@ def parallel_read_into_storage( work.put(None) size = os.path.getsize(path) - use_direct = ( - _env_flag("FLASHPACK_DIRECT_IO") - and hasattr(os, "O_DIRECT") - # A page-cache-hot file is faster through buffered reads; O_DIRECT - # would bypass the cache and re-fetch from the filesystem. - and _page_cache_resident_fraction(path, size) < 0.9 - ) + # A page-cache-hot file is faster through buffered reads; O_DIRECT + # would bypass the cache and re-fetch from the filesystem. + use_direct = _should_use_direct(path, size) pool = _get_pinned_pool(n_threads, chunk_bytes) errors: list[BaseException] = [] diff --git a/tests/test_sample_probe.py b/tests/test_sample_probe.py new file mode 100644 index 0000000..ee981f7 --- /dev/null +++ b/tests/test_sample_probe.py @@ -0,0 +1,87 @@ +"""Unit tests for the O_DIRECT sample-read gate (_should_use_direct).""" + +import os + +import pytest + + +class TestSampleProbe: + def test_small_file_reports_hot(self, tmp_path) -> None: + from flashpack.parallel_read import _sample_read_gbps + + p = tmp_path / "small.bin" + p.write_bytes(b"x" * 1024) + assert _sample_read_gbps(str(p), 1024) == float("inf") + + @pytest.mark.skipif( + not hasattr(os, "preadv"), + reason="the sample probe reads with preadv (POSIX-only); on platforms " + "without it the probe reports cold and the reader keeps O_DIRECT off", + ) + def test_probe_returns_rate_on_real_file(self, tmp_path) -> None: + from flashpack.parallel_read import ( + _SAMPLE_PROBE_BYTES, + _SAMPLE_PROBE_THREADS, + _sample_read_gbps, + ) + + size = _SAMPLE_PROBE_THREADS * _SAMPLE_PROBE_BYTES + (1 << 20) + p = tmp_path / "big.bin" + with open(p, "wb") as f: + f.truncate(size) + rate = _sample_read_gbps(str(p), size) + assert rate > 0.0 and rate != float("inf") + + def test_direct_io_env_kill_switch(self, tmp_path, monkeypatch) -> None: + from flashpack.parallel_read import _should_use_direct + + monkeypatch.setenv("FLASHPACK_DIRECT_IO", "0") + p = tmp_path / "f.bin" + p.write_bytes(b"x" * 4096) + assert _should_use_direct(str(p), 4096) is False + + def test_probe_disabled_falls_back_to_mincore_only( + self, tmp_path, monkeypatch + ) -> None: + import flashpack.parallel_read as pr + + monkeypatch.setenv("FLASHPACK_SAMPLE_PROBE", "0") + monkeypatch.delenv("FLASHPACK_DIRECT_IO", raising=False) + monkeypatch.setattr(pr, "_page_cache_resident_fraction", lambda *_: 0.0) + monkeypatch.setattr( + pr, + "_sample_read_gbps", + lambda *_: pytest.fail("probe must not run when disabled"), + ) + p = tmp_path / "f.bin" + p.write_bytes(b"x" * 4096) + if not hasattr(os, "O_DIRECT"): + pytest.skip("no O_DIRECT on this platform") + assert pr._should_use_direct(str(p), 4096) is True + + def test_hot_probe_blocks_direct(self, tmp_path, monkeypatch) -> None: + import flashpack.parallel_read as pr + + monkeypatch.delenv("FLASHPACK_DIRECT_IO", raising=False) + monkeypatch.delenv("FLASHPACK_SAMPLE_PROBE", raising=False) + monkeypatch.setattr(pr, "_page_cache_resident_fraction", lambda *_: 0.0) + monkeypatch.setattr(pr, "_sample_read_gbps", lambda *_: 21.6) + p = tmp_path / "f.bin" + p.write_bytes(b"x" * 4096) + if not hasattr(os, "O_DIRECT"): + pytest.skip("no O_DIRECT on this platform") + assert pr._should_use_direct(str(p), 4096) is False + + def test_mincore_hot_short_circuits_probe(self, tmp_path, monkeypatch) -> None: + import flashpack.parallel_read as pr + + monkeypatch.delenv("FLASHPACK_DIRECT_IO", raising=False) + monkeypatch.setattr(pr, "_page_cache_resident_fraction", lambda *_: 1.0) + monkeypatch.setattr( + pr, + "_sample_read_gbps", + lambda *_: pytest.fail("probe must not run when mincore says hot"), + ) + p = tmp_path / "f.bin" + p.write_bytes(b"x" * 4096) + assert pr._should_use_direct(str(p), 4096) is False