fix(layer-stream): align a pinned staging arena whose base is off a sector boundary (#1531)

changelog.d/0.75.1/1621.fixed.md

@@ -0,0 +1,11 @@
+- **A pinned staging arena whose base is off a 4096-byte boundary no longer
+  turns direct I/O off for the whole disk-tier store (#1531 by @ybai08 in
+  #1621).** `plan_pinned_arenas` aligns the staging regions relative to their
+  arena, and torch's caching host allocator can hand back an arena at a
+  multiple of only 512 bytes; one misaligned region then sent every layer
+  through the page cache, 2.4x slower cold (#974), with a `logger.warning`
+  as the only sign. `AsyncDiskSource` now reads each arena's base, shifts the
+  placements inside the arena when its slack allows, and otherwise takes the
+  next power of two up; `pinned_bytes` counts the arenas actually held. If the
+  staging is misaligned anyway, the page-cache fallback is printed on the
+  console as well as logged.

src/soup_cli/utils/async_disk_source.py

@@ -64,6 +64,11 @@
 
 logger = logging.getLogger(__name__)
 
+UNALIGNED_STAGING_MESSAGE = (
+    "layer streaming's staging is not sector-aligned on this box; the "
+    "disk tier reads through the page cache instead of direct I/O"
+)
+
 #: How many byte ranges a layer's data section is read as, each by its own
 #: worker thread through its own direct-I/O handle. Measured cold on the
 #: 70B-shaped NF4 store (#974; fresh layers per configuration, Samsung PM9B1
@@ -330,6 +335,28 @@ def done(self) -> bool:
         return self.outcomes is not None and all(o.done.is_set() for o in self.outcomes)
 
 
+def _allocate_pinned_arena(size: int) -> Any:
+    """``size`` bytes of page-locked CPU memory, or an error that says why not."""
+    import torch
+
+    arena = torch.empty(size, dtype=torch.uint8, device="cpu", pin_memory=True)
+    if arena.device.type != "cpu":
+        raise RuntimeError(
+            "layer streaming's async disk source requested a CPU "
+            f"tensor, but torch returned {arena.device}."
+        )
+    # Mirrors RamSource: a box that hands back pageable memory
+    # would otherwise report pinned=True while silently paying
+    # the ~97% -> ~79% GPU-utilisation cost of a synchronous
+    # host-to-device copy.
+    if not arena.is_pinned():
+        raise RuntimeError(
+            "layer streaming requested pinned CPU RAM, but torch "
+            "returned pageable memory; retry with pin=False."
+        )
+    return arena
+
+
 class AsyncDiskSource:
     """Read layers ahead on a background thread; ``get`` hands over the result.
 
@@ -557,11 +584,11 @@ def __init__(
         for idx, root in enumerate(self._root_of):
             first_of_root.setdefault(root, idx)
         aligned = all(region.data_ptr() % SECTOR_BYTES == 0 for region in self._regions)
-        if not aligned:  # pragma: no cover — cudaHostAlloc hands out page-aligned blocks
-            logger.warning(
-                "layer streaming's staging is not sector-aligned on this box; the "
-                "disk tier reads through the page cache instead of direct I/O"
-            )
+        #: False when the whole source reads through the page cache for this reason;
+        #: ``_build_source`` says so on the console.
+        self.staging_aligned = aligned
+        if not aligned:  # _allocate_staging aligns every arena; kept as the safe fallback
+            logger.warning(UNALIGNED_STAGING_MESSAGE)
         self._open_of: Dict[int, Any] = {}
         for root in self._roots:
             self._open_of[root] = self._open_buffered
@@ -680,26 +707,26 @@ def _allocate_staging(self, region_sizes: Sequence[int]) -> Tuple[List[Any], Lis
         regions: List[Any] = []
         if self.pinned:
             plan = plan_pinned_arenas(region_sizes, align=SECTOR_BYTES)
-            for size in plan.arena_sizes:
-                arena = torch.empty(size, dtype=torch.uint8, device="cpu", pin_memory=True)
-                if arena.device.type != "cpu":
-                    raise RuntimeError(
-                        "layer streaming's async disk source requested a CPU "
-                        f"tensor, but torch returned {arena.device}."
-                    )
-                # Mirrors RamSource: a box that hands back pageable memory
-                # would otherwise report pinned=True while silently paying
-                # the ~97% -> ~79% GPU-utilisation cost of a synchronous
-                # host-to-device copy.
-                if not arena.is_pinned():
-                    raise RuntimeError(
-                        "layer streaming requested pinned CPU RAM, but torch "
-                        "returned pageable memory; retry with pin=False."
-                    )
+            fills = [0] * len(plan.arena_sizes)
+            for (arena_index, offset), size in zip(plan.placements, region_sizes):
+                fills[arena_index] = max(fills[arena_index], offset + size)
+            # The placements are sector-aligned within their arena, but the caching
+            # host allocator can hand back a base at a multiple of only 512 bytes.
+            shifts: List[int] = []
+            for size, fill in zip(plan.arena_sizes, fills):
+                arena = _allocate_pinned_arena(size)
+                shift = (-arena.data_ptr()) % SECTOR_BYTES
+                if shift + fill > size:
+                    # No room to shift: the next power of two has at least a sector spare.
+                    del arena
+                    arena = _allocate_pinned_arena(2 * size)
+                    shift = (-arena.data_ptr()) % SECTOR_BYTES
                 arenas.append(arena)
+                shifts.append(shift)
             for (arena_index, offset), size in zip(plan.placements, region_sizes):
-                regions.append(arenas[arena_index][offset : offset + size])
-            self.pinned_bytes = plan.pinned_bytes
+                start = shifts[arena_index] + offset
+                regions.append(arenas[arena_index][start : start + size])
+            self.pinned_bytes = sum(arena.numel() for arena in arenas)
         else:
             # The CPU allocator neither rounds nor aligns: a sector of slack
             # per region buys the alignment direct I/O needs.

src/soup_cli/utils/layer_stream_runtime.py

@@ -3222,6 +3222,17 @@ def _require_stripe_folders(shard_dir: str, index: Any) -> None:
             )
 
 
+def _say_unaligned_staging(source: Any, console: Any) -> None:
+    """Put the disk tier's page-cache fallback where a `soup train` user reads it (#1531).
+
+    The source logs it; the logger is not what the Rich output shows.
+    """
+    if console is not None and not getattr(source, "staging_aligned", True):
+        from soup_cli.utils.async_disk_source import UNALIGNED_STAGING_MESSAGE
+
+        console.print(f"[yellow]{UNALIGNED_STAGING_MESSAGE}[/]")
+
+
 def _build_source(
     shard_dir,
     n_layers,
@@ -3323,8 +3334,10 @@ def _build_source(
                 # otherwise report as an out-of-memory it never had.
                 recover_from_failed_page_lock(console=console)
             else:
+                _say_unaligned_staging(source, console)
                 return source, source.pinned
         source = AsyncDiskSource(shard_dir, n_layers, spec, pin=False, **open_kwargs)
+        _say_unaligned_staging(source, console)
         return source, source.pinned
     # `source.pinned` on both branches rather than the literal, so the tuple's
     # second element has ONE meaning to read off: what the source says about

tests/test_issue1531_pinned_arena_alignment.py

@@ -0,0 +1,225 @@
+"""#1531 — a pinned staging arena whose base is off a 4096-byte boundary no longer turns
+direct I/O off for the whole store.
+
+``plan_pinned_arenas`` keeps every placement on a sector boundary RELATIVE to its
+arena, so the regions are aligned only when the arena's base is. torch's caching host
+allocator does not promise that: on torch 2.14.0+cu130 a probe of 400 pinned
+allocations with frees in between got 159 bases at a multiple of 512 but not of 4096.
+The source now aligns inside the arena when its slack allows, and otherwise takes the
+next power of two up, which always has room for the shift.
+
+No GPU here: the arena allocation is replaced with pageable CPU memory at a forced
+512-byte offset, the shape the probe saw.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("safetensors")
+
+import soup_cli.utils.async_disk_source as mod  # noqa: E402
+from soup_cli.utils.async_disk_source import AsyncDiskSource  # noqa: E402
+from tests.test_issue974_direct_range_reads import N_LAYERS, _shards, _spec  # noqa: E402
+
+SECTOR = 4096
+OFF_BY = 512
+
+
+def _misaligned_arena(size: int, off_by: int = OFF_BY):
+    """``size`` bytes whose first byte sits ``off_by`` past a sector boundary."""
+    buffer = torch.empty(size + 2 * SECTOR, dtype=torch.uint8)
+    pad = (off_by - buffer.data_ptr()) % SECTOR
+    return buffer[pad : pad + size]
+
+
+@pytest.fixture
+def misaligned(monkeypatch):
+    sizes = []
+
+    def allocate(size: int):
+        sizes.append(size)
+        return _misaligned_arena(size)
+
+    monkeypatch.setattr(mod, "_allocate_pinned_arena", allocate)
+    return sizes
+
+
+def _bare_pinned_source() -> AsyncDiskSource:
+    source = AsyncDiskSource.__new__(AsyncDiskSource)
+    source.pinned = True
+    return source
+
+
+class TestAllocateStaging:
+    def test_the_fake_arena_is_off_by_512(self):
+        assert _misaligned_arena(16 * SECTOR).data_ptr() % SECTOR == OFF_BY
+
+    def test_a_misaligned_base_with_slack_is_aligned_inside_the_same_arena(self, misaligned):
+        # Three sectors pack into a 16 KiB arena: one sector of slack covers the shift.
+        arenas, regions = _bare_pinned_source()._allocate_staging([3 * SECTOR])
+        assert misaligned == [4 * SECTOR]
+        assert [a.numel() for a in arenas] == [4 * SECTOR]
+        assert all(r.data_ptr() % SECTOR == 0 for r in regions)
+        assert [r.numel() for r in regions] == [3 * SECTOR]
+
+    def test_a_full_misaligned_arena_is_replaced_by_the_next_size_up(self, misaligned):
+        # Two regions fill a 32 KiB arena exactly: no room to shift inside it.
+        source = _bare_pinned_source()
+        arenas, regions = source._allocate_staging([4 * SECTOR, 4 * SECTOR])
+        assert misaligned == [8 * SECTOR, 16 * SECTOR]
+        assert [a.numel() for a in arenas] == [16 * SECTOR]
+        assert all(r.data_ptr() % SECTOR == 0 for r in regions)
+        assert [r.numel() for r in regions] == [4 * SECTOR, 4 * SECTOR]
+        # The page-locked total counts the arena actually held, not the plan's.
+        assert source.pinned_bytes == 16 * SECTOR
+
+    def test_the_shift_is_measured_again_on_the_bigger_arena(self, monkeypatch):
+        # The replacement arena comes back at a different offset from the first.
+        offsets = iter([OFF_BY, 2 * OFF_BY])
+        monkeypatch.setattr(
+            mod, "_allocate_pinned_arena", lambda size: _misaligned_arena(size, next(offsets))
+        )
+        arenas, regions = _bare_pinned_source()._allocate_staging([4 * SECTOR, 4 * SECTOR])
+        assert arenas[0].data_ptr() % SECTOR == 2 * OFF_BY
+        assert all(r.data_ptr() % SECTOR == 0 for r in regions)
+
+    def test_the_regions_of_one_arena_do_not_overlap_after_the_shift(self, misaligned):
+        arenas, regions = _bare_pinned_source()._allocate_staging([SECTOR, 2 * SECTOR, SECTOR])
+        assert len(arenas) == 1
+        spans = sorted((r.data_ptr(), r.data_ptr() + r.numel()) for r in regions)
+        assert all(end <= start for (_, end), (start, _) in zip(spans, spans[1:]))
+        lo, hi = arenas[0].data_ptr(), arenas[0].data_ptr() + arenas[0].numel()
+        assert all(lo <= start and end <= hi for start, end in spans)
+
+    def test_an_aligned_base_is_used_as_planned(self, monkeypatch):
+        sizes = []
+
+        def aligned(size: int):
+            sizes.append(size)
+            buffer = torch.empty(size + SECTOR, dtype=torch.uint8)
+            pad = (-buffer.data_ptr()) % SECTOR
+            return buffer[pad : pad + size]
+
+        monkeypatch.setattr(mod, "_allocate_pinned_arena", aligned)
+        source = _bare_pinned_source()
+        arenas, regions = source._allocate_staging([4 * SECTOR, 4 * SECTOR])
+        assert sizes == [8 * SECTOR]
+        assert regions[0].data_ptr() == arenas[0].data_ptr()
+        assert source.pinned_bytes == 8 * SECTOR
+
+
+class TestTheSource:
+    def test_a_misaligned_arena_keeps_the_regions_aligned_and_says_nothing(
+        self, tmp_path, misaligned, caplog
+    ):
+        shard_dir = _shards(tmp_path)
+        with caplog.at_level(logging.WARNING, logger=mod.__name__):
+            source = AsyncDiskSource(shard_dir, N_LAYERS, _spec(shard_dir), read_ahead=2, pin=True)
+        try:
+            assert misaligned, "the pinned arenas did not go through the patched allocation"
+            assert all(r.data_ptr() % SECTOR == 0 for r in source._regions)
+            assert not [r for r in caplog.records if "not sector-aligned" in r.getMessage()]
+            assert source.pinned_bytes == sum(a.numel() for a in source._arenas)
+        finally:
+            source.close()
+
+    def test_misaligned_regions_still_fall_back_to_the_page_cache(
+        self, tmp_path, monkeypatch, caplog
+    ):
+        """The fallback stays for staging that is misaligned anyway: reading it with
+        direct I/O would fail, so the whole source reads through the page cache."""
+        real = AsyncDiskSource._allocate_staging
+
+        def shifted(self, region_sizes):
+            arenas, regions = real(self, region_sizes)
+            return arenas, [_misaligned_arena(r.numel()) for r in regions]
+
+        monkeypatch.setattr(AsyncDiskSource, "_allocate_staging", shifted)
+        opened = []
+        monkeypatch.setattr(mod, "open_direct", lambda path: opened.append(path))
+        shard_dir = _shards(tmp_path)
+        with caplog.at_level(logging.WARNING, logger=mod.__name__):
+            source = AsyncDiskSource(shard_dir, N_LAYERS, _spec(shard_dir), read_ahead=2, pin=False)
+        try:
+            assert source.direct_io is False
+            assert opened == [], "direct I/O was probed for misaligned staging"
+            assert [r for r in caplog.records if "not sector-aligned" in r.getMessage()]
+        finally:
+            source.close()
+
+
+class _Console:
+    def __init__(self):
+        self.lines = []
+
+    def print(self, text):
+        self.lines.append(text)
+
+
+@pytest.fixture
+def unalignable(monkeypatch):
+    """Staging that is misaligned whatever the arenas were: the fallback's one case."""
+    real = AsyncDiskSource._allocate_staging
+
+    def shifted(self, region_sizes):
+        arenas, regions = real(self, region_sizes)
+        return arenas, [_misaligned_arena(r.numel()) for r in regions]
+
+    monkeypatch.setattr(AsyncDiskSource, "_allocate_staging", shifted)
+
+
+class TestTheConsoleLine:
+    """A `soup train` user sees the Rich console, not the logger."""
+
+    def _build(self, tmp_path, console, *, pin):
+        from soup_cli.utils.layer_stream_runtime import _build_source
+
+        shard_dir = _shards(tmp_path)
+        return _build_source(
+            shard_dir, N_LAYERS, _spec(shard_dir), pin, console, tier="disk", read_ahead=2
+        )
+
+    @staticmethod
+    def _fallback_lines(console):
+        return [line for line in console.lines if "page cache" in line]
+
+    def test_a_page_cache_fallback_is_said_on_the_console(self, tmp_path, unalignable):
+        console = _Console()
+        source, pinned = self._build(tmp_path, console, pin=False)
+        try:
+            assert source.direct_io is False and pinned is False
+            lines = self._fallback_lines(console)
+            assert len(lines) == 1
+            assert "not sector-aligned" in lines[0] and "direct I/O" in lines[0]
+        finally:
+            source.close()
+
+    def test_a_pinned_source_says_it_too(self, tmp_path, misaligned, unalignable):
+        console = _Console()
+        source, pinned = self._build(tmp_path, console, pin=True)
+        try:
+            assert source.direct_io is False and pinned is True
+            assert len(self._fallback_lines(console)) == 1
+        finally:
+            source.close()
+
+    def test_aligned_staging_says_nothing(self, tmp_path):
+        console = _Console()
+        source, _ = self._build(tmp_path, console, pin=False)
+        try:
+            assert source.staging_aligned is True
+            assert self._fallback_lines(console) == []
+        finally:
+            source.close()
+
+    def test_without_a_console_the_logger_says_it_once(self, tmp_path, unalignable, caplog):
+        with caplog.at_level(logging.WARNING, logger=mod.__name__):
+            source, _ = self._build(tmp_path, None, pin=False)
+        try:
+            assert len([r for r in caplog.records if "not sector-aligned" in r.getMessage()]) == 1
+        finally:
+            source.close()