Repository navigation
[rocjitsu] Reduce memory keeper startup and topology IPC - #14999
Conversation
github-actions
Bot
added
documentation
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
There was a problem hiding this comment.
🟡 Changes recommended
Fork-inherited sockets can leave keeper channels allocated after probes or owner exit, while interrupted health checks can prevent retrying against a replacement keeper.
3 open issues
What changed in this PR
This PR reduces rocjitsu memory-keeper startup and topology IPC overhead following #14039, while preserving isolation from application file limits and compatibility with exported tokens.
Changes:
- Batch nonempty topology files, inline payloads up to 4 KiB, and acknowledge up to eight writes together.
- Prepare the keeper during interposer initialization and use private connections for allocation replies.
- Separate the keeper server and wire protocol, with expanded lifecycle, resource-limit, and recovery tests.
| File | Description |
|---|---|
emulation/rocjitsu/tools/rocjitsu_memory.cpp |
Pass the listener name to the keeper entry point. |
emulation/rocjitsu/tools/CMakeLists.txt |
Link the helper to server and protocol sources. |
emulation/rocjitsu/tests/memory_backing_test.cpp |
Test batching, recovery, limits, and channel lifetimes. |
emulation/rocjitsu/tests/interposer_dup_test.cpp |
Test preparation before application file limits drop. |
emulation/rocjitsu/tests/CMakeLists.txt |
Register the new interposer test. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/sysfs.h |
Add pending topology-file storage. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/sysfs.cpp |
Queue topology files for batched population. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_keeper.h |
Define shared protocol types and startup handshake. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_keeper.cpp |
Extract the server and implement batched writes. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_keeper_protocol.cpp |
Centralize descriptor transport and protocol validation. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_backing.h |
Expose keeper preparation and batch-file APIs. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_backing.cpp |
Implement private requests, batching, and recovery. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/libc_passthrough.h |
Share host-call and descriptor-close helpers. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/interposer.cpp |
Replace the warmup allocation with keeper preparation. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/device_backing.h |
Extract shared device-backing identifiers. |
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/CMakeLists.txt |
Include protocol transport in the Linux backend. |
emulation/rocjitsu/docs/memory-backing.md |
Document preparation, batching, and connection behavior. |
Review effort: Balanced
e8166ef to
cf97376
Compare
There was a problem hiding this comment.
🔵 Needs a closer look
The cross-process lifecycle and IPC refactor requires Linux integration and sanitizer testing to validate concurrent fork/exec, descriptor-pressure handling, and helper-loss recovery beyond source inspection.
No open issues 🎉
3 resolved since last review
Review effort: Balanced
sjain-stanford
left a comment
sjain-stanford
left a comment
There was a problem hiding this comment.
Could we add each revision's range or median absolute deviation and a cycles:u/instructions:u cross-check for the small wall-time improvement, as requested by the sampling protocol in emulation/rocjitsu/docs/benchmarking.md? The 80 pairs and median deltas do not show the spread on this unreserved host. If counters are unavailable, an identical-binary control would help establish the noise floor.
Please also label the syscall counts with the measured process boundary. I reproduced 1 socketpair and 9 recvmsg calls in the application process at this head, alongside 40 sendmsg calls; the helper separately receives the file requests. Unless a separate measurement includes the keeper, “the application process” would be more accurate than “the process tree.”
We should use descriptive indices and explicit view/reference types in the new batch code, following emulation/rocjitsu/docs/style.md:
- In
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_backing.cpp:697, use a name such asfile_indexfor the acknowledgment index. Inemulation/rocjitsu/tests/memory_backing_test.cpp:490, usefile_indexandbyte_indexfor the nested payload-generation loops; these distinguish the file number from the byte offset in the formula. - In
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_backing.cpp:801and:813, spell outBackingFile &firstandconst std::span<BackingFile> batch. The latter makes the borrowed view and mutable entries explicit when the following code updates the caller's errors.
I think we should make the nonempty-run contract explicit in emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_backing.h:70 and explain the ordering at emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/sysfs.cpp:529. The header promises one connection for the call, but create_backing_files closes it at each empty entry and truncates that entry locally; docs/memory-backing.md already describes this behavior. Keeping version before the empty dev_dri files lets the generated tree use one connection.
A useful optional regression guard would prepare the keeper, count private connect calls during Sysfs::generate for a representative multi-GPU tree, and assert one. Counting FlushFiles requests would also protect the grouping of eight inline writes. Existing content checks do not catch an extra connection caused by changing the write order.
e73876b to
994c307
Compare
sjain-stanford
left a comment
sjain-stanford
left a comment
There was a problem hiding this comment.
We should use descriptive indices and explicit view/reference types in the new batch code, following
emulation/rocjitsu/docs/style.md:
- In
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_backing.cpp:697, use a name such asfile_indexfor the acknowledgment index. Inemulation/rocjitsu/tests/memory_backing_test.cpp:490, usefile_indexandbyte_indexfor the nested payload-generation loops; these distinguish the file number from the byte offset in the formula.- In
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_backing.cpp:801and:813, spell outBackingFile &firstandconst std::span<BackingFile> batch. The latter makes the borrowed view and mutable entries explicit when the following code updates the caller's errors.
We should apply the same descriptive-name and explicit-type conventions to the remaining new batch code, following emulation/rocjitsu/docs/style.md:
- In
emulation/rocjitsu/tests/memory_backing_test.cpp:1656and:1670, usefile_index; at:387, usebyte_index. These names make the acknowledged-file boundary and payload byte offset explicit, as incheck_metadata_batch. - In
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/memory_keeper.cpp:344, spell outKeeperRequest &request. The member initializer does not show its type; naming the wire header makes the subsequent size and operation checks easier to read without looking upKeeperFilePacket.
sjain-stanford
left a comment
sjain-stanford
left a comment
There was a problem hiding this comment.
LGTM modulo two comments.
Populate runs of nonempty topology files over temporary connections. Send files of up to 4 KiB inline and acknowledge up to eight writes at a time. Drain pending writes at the descriptor limit and keep acknowledged files if the helper restarts. Preserve per-file open errors and ordered creation and truncation if a connection fails. Keep mapped writes for larger files and caller-owned snapshots. Link rocjitsu-memory from its server loop and the shared wire protocol. Create the listener after exec so application children cannot retain it. Wait for the ready reply during interposer initialization so the keeper inherits the original file limits. This removes the throwaway backing allocation. Tie each private connection to the helper's pidfd inode or require the cookie it sent over the inherited startup channel. Bound connections to a full listener and retain the original helper if it cannot answer under pressure. Keep the existing exported token protocol. Compare the pre-rebase build 994c307 with develop at 7215ce9 on a Threadripper PRO 9995WX using GCC 16 Release, x86-64-v4, a fixed 2.5 GHz core, ASLR disabled and an exclusive smoke lock. Alternate baseline and candidate order across 80 unreserved pairs. Collect userspace cycles and instructions in a separate 80-pair run, including the application and its descendants through keeper exit. For gfx1250 rocminfo, median time falls from 54.581 ms (MAD 0.530) to 53.802 ms (MAD 0.345). Paired time improves 1.39% (0.762 ms, 95% CI 0.88-1.90%). Paired userspace cycles change by +0.16% (95% CI -0.34 to +0.56%) and instructions by -0.16%. For direct preload, median time falls from 4.875 ms (MAD 0.058) to 4.744 ms (MAD 0.060). Paired time improves 3.53% (0.171 ms, 95% CI 2.47-4.15%), with 3.86% fewer cycles and 1.90% fewer instructions. In the rocminfo application process, socketpair calls fall from 38 to 1, recvmsg from 38 to 9, and sendmsg from 69 to 40. Five connections check peer credentials and the pidfd inode. The helper handles the file requests separately. Its binary shrinks from 2,136,408 bytes and 2,990 relative relocations to 35,864 bytes and three. Follow-up to the startup performance regression introduced in #14039.
994c307 to
6333a3b
Compare
rocjitsu post-commit fuzz passedThe 32-worker multiplexer campaigns for commit
|
TheRock Submodule Bump ActivityNewest first
|
Populate runs of nonempty topology files over temporary connections.
Send files of up to 4 KiB inline and acknowledge up to eight writes at a
time. Drain pending writes at the descriptor limit and keep acknowledged
files if the helper restarts. Preserve per-file open errors and ordered
creation and truncation if a connection fails. Keep mapped writes for
larger files and caller-owned snapshots.
Link rocjitsu-memory from its server loop and the shared wire protocol.
Create the listener after exec so application children cannot retain it.
Wait for the ready reply during interposer initialization so the keeper
inherits the original file limits. This removes the throwaway backing
allocation. Tie each private connection to the helper's pidfd inode or
require the cookie it sent over the inherited startup channel. Bound
connections to a full listener and retain the original helper if it
cannot answer under pressure. Keep the existing exported token protocol.
Compare the pre-rebase build 994c307 with develop at 7215ce9
on a Threadripper PRO 9995WX using GCC 16 Release, x86-64-v4, a fixed
2.5 GHz core, ASLR disabled and an exclusive smoke lock. Alternate
baseline and candidate order across 80 unreserved pairs.
Collect userspace cycles and instructions in a separate 80-pair run,
including the application and its descendants through keeper exit.
For gfx1250 rocminfo, median time falls from 54.581 ms (MAD 0.530)
to 53.802 ms (MAD 0.345). Paired time improves 1.39% (0.762 ms,
95% CI 0.88-1.90%). Paired userspace cycles change by +0.16% (95% CI
-0.34 to +0.56%) and instructions by -0.16%. For direct preload,
median time falls from 4.875 ms (MAD 0.058) to 4.744 ms (MAD 0.060).
Paired time improves 3.53% (0.171 ms, 95% CI 2.47-4.15%), with 3.86%
fewer cycles and 1.90% fewer instructions.
In the rocminfo application process, socketpair calls fall from 38 to
1, recvmsg from 38 to 9, and sendmsg from 69 to 40. Five connections
check peer credentials and the pidfd inode. The helper handles the file
requests separately. Its binary shrinks from 2,136,408 bytes and 2,990
relative relocations to 35,864 bytes and three.
Follow-up to the startup performance regression introduced in #14039.