Skip to content

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

Merged
MakazhanAlpamys merged 2 commits into
MakazhanAlpamys:mainfrom
ybai08:fix/1531-align-pinned-arena-base
Oct 6, 2026
Merged

MakazhanAlpamys merged 2 commits into
MakazhanAlpamys:mainfrom
ybai08:fix/1531-align-pinned-arena-base

Conversation

@ybai08

@ybai08 ybai08 commented Oct 4, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Closes #1531, on current main (after #1536), in two commits. The second is item 3 on its own, as asked in the issue.

1. A misaligned arena base no longer turns direct I/O off (_allocate_staging)

plan_pinned_arenas keeps the placements on sector boundaries relative to their arena, so the regions are aligned only when the arena's base is. The base is now read right after the arena is allocated:

  • If it is off a 4096-byte boundary and the packed regions still fit, the placements are shifted by (-base) % SECTOR_BYTES inside the same arena.
  • If they do not fit, the arena is dropped and one of the next power of two is taken. That one always has room, and its shift is measured again.
  • pinned_bytes counts the arenas that are held, so a bigger arena shows there.

The allocation is a module-level function, _allocate_pinned_arena(size), so a CPU test can replace it. The pragma: no cover on the page-cache branch is gone.

2. The fallback is said on the console (item 3)

When the staging is still not sector-aligned, the source read the whole store through the page cache and said so only in logger.warning. The source now records staging_aligned, and _build_source prints the same line on the console at both places it returns a disk source. direct_io, which stream_probe.py records, is already False in that case after #1536. I did not put the per-drive "direct I/O is unavailable there" message from #1536 on the console: that is a property of the volume, not this fallback. Tell me if you want it there too.

Tests (tests/test_issue1531_pinned_arena_alignment.py, 12 tests, CPU only: the arena allocation is replaced by pageable memory at a forced 512-byte offset)

  • The shift fits: same arena, regions aligned.
  • The arena is full: replaced by the next size up, regions aligned, pinned_bytes is the bigger arena.
  • The replacement arena comes back at a different offset: the shift is measured again.
  • The regions of one arena do not overlap after the shift and stay inside it.
  • An aligned base is used as planned.
  • Through the real constructor: a misaligned arena gives aligned regions and no warning.
  • Staging that is misaligned anyway: direct_io is False, direct I/O is not probed, the logger says it (the former pragma: no cover branch).
  • The console: said once for a pageable source and once for a pinned one, nothing for aligned staging, and the logger alone says it when there is no console.

Each guard was removed in turn: 12 mutations, 12 caught (6 in _allocate_staging, 6 in the console path). One survived my first version of the tests, the re-measured shift, which is why that test is there.

What I ran

  • pytest tests/test_issue1531_pinned_arena_alignment.py --no-cov: 12 passed; 8 passed at the first commit alone.
  • The 18 test files that mention _build_source or AsyncDiskSource: 901 passed, 45 skipped. Six tests in test_v07200.py and test_disk_tier_stripe_read_ahead_advice.py build the source from a stand-in class, so _build_source reads staging_aligned with a default of aligned.
  • pytest tests/ -q --no-cov: 32866 passed, 254 skipped, 2 failed. That run was on the head before the re-measured-shift test was added (a test-only change). The two failures are tests/test_v0537.py::TestToolEndpointsLive::test_bash_tool_bounded_streaming_kills_massive_output and ::test_bash_tool_environment_secret_isolation. They fail the same way on main on this machine: the sandboxed python3 exits 72 because of the state of this Mac's Xcode tools.
  • ruff check on the three changed files: clean.
  • Not run: the third acceptance criterion, the full run of the layer-streaming test files on a CUDA box. This machine is Apple silicon with no CUDA, so that run is yours, as discussed in the issue.

Type of change

  • Bug fix
  • New feature
  • Refactor / cleanup
  • Documentation
  • Tests

Checklist

  • ruff check src/soup_cli/ scripts/ tests/ passes (run on the changed files)
  • pytest tests/ -v passes (2 failures that also fail on main on this machine, see above)
  • Updated relevant docs (README.md and the matching page under docs/) if needed (the direct I/O paragraph in docs/performance-and-quantization.md does not describe this fallback)
  • Added a changelog fragment, or this change has no user-visible impact

ybai08 added 2 commits October 4, 2026 11:56
…ector boundary (MakazhanAlpamys#1531)

plan_pinned_arenas keeps placements sector-aligned within their arena, but
torch's caching host allocator can hand back a base at a multiple of only 512
bytes, and one misaligned region turned direct I/O off for the whole store.
Shift the arena's placements by the base's distance to the next sector when
its slack allows, else take the next power of two, which always has room;
pinned_bytes counts the arenas actually held.
…sole (MakazhanAlpamys#1531)

When the staging is not sector-aligned the source reads the whole store
through the page cache and said so only in the logger, which a soup train
user does not see. The source now records staging_aligned and _build_source
prints the same line on the console.

@MakazhanAlpamys MakazhanAlpamys left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ybai08, thank you, this is a careful fix. Two things stood out. You turned the arena allocation into a small seam (_allocate_pinned_arena), so a CPU test can force the 512-byte offset without faking torch itself. And you measure the shift again on the replacement arena instead of reusing the first one. test_the_shift_is_measured_again_on_the_bigger_arena is the only test that catches the loss of that line, so it earns its place. Keeping item 3 as its own commit on top of #1536 made the review easy.

What I checked on 413e3510 (CPU only):

  • The next size always has room. plan_pinned_arenas floors every arena at align, so a staging arena is at least 4096 bytes and 2 * size - fill >= size, which is more than any shift. direct_io is False whenever the staging is unaligned, because every drive keeps the buffered opener.
  • Your new file passes: 12 passed. The same file on current main gives 4 failed, 5 errors, 3 passed. The three that pass there are the helper check, the former pragma: no cover branch and the logger-only path, as expected.
  • A trial merge with current main is clean. On the merged tree, 16 test files that build or patch the disk source give 860 passed, 47 skipped, and ruff check is clean.
  • 12 mutations, all caught. Seven are in _allocate_staging: the first shift forced to 0, the next-size branch disabled, the shift not measured again, pinned_bytes back to the plan's figure, > turned into >=, the replacement at size instead of 2 * size, and the shift ignored when slicing. One forces staging_aligned to True. Four are on the console path, among them the getattr default flipped, which test_v07200.py catches.

To your question: please keep the per-drive message out of this PR. #1615 covers it, together with the ready line and stats(), and it now has its own contributor.

The CUDA run of the layer-streaming test files (the issue's third criterion) stays with us, as agreed on the issue. From my side this is ready: it now waits only on the required checks (the full matrix starts when it carries ci:full).

@MakazhanAlpamys MakazhanAlpamys added the ci:full Run the full CI matrix on this PR (a maintainer adds it at approval) label Oct 5, 2026
@MakazhanAlpamys

Copy link
Copy Markdown
Owner

Closing and reopening only to make CI test this against today's main (a dependency release broke CI repo-wide, fixed by #1637). Nothing for you to do; the approval stands.

@MakazhanAlpamys MakazhanAlpamys added ci:full Run the full CI matrix on this PR (a maintainer adds it at approval) and removed ci:full Run the full CI matrix on this PR (a maintainer adds it at approval) labels Oct 6, 2026
@MakazhanAlpamys
MakazhanAlpamys merged commit f87423d into MakazhanAlpamys:main Oct 6, 2026
33 of 57 checks passed
@MakazhanAlpamys

Copy link
Copy Markdown
Owner

Merged as f87423d6. Thank you, @ybai08: you turned the arena allocation into a small seam so a CPU test can force the 512-byte offset, and you measure the shift again on the replacement arena, which one of your tests alone pins. I ran your 12 tests (all pass; on main without the fix, 4 failed and 5 errored), the 16 disk-source test files on the merged tree (860 passed, 47 skipped) and 12 mutations of the fix and the console path, all 12 caught.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci:full Run the full CI matrix on this PR (a maintainer adds it at approval)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

layer streaming: a pinned staging arena can start off a 4096-byte boundary, which silently turns direct I/O off

2 participants