wp: harden io_uring submit accounting

Root cause: the host and P2P FileIOLayer io_uring paths counted a request as pending when an SQE was prepared, then ignored the return from io_uring_submit. Under resident-dense paging, many small expert reads and heavy eviction can fill the CQ/SQ pressure window at depth 8/16. If io_uring_submit returned -EBUSY/-EAGAIN/0 or otherwise submitted fewer entries than prepared, the pager still waited for req_ids that were not actually kernel-owned. The targeted waiter then blocked in io_cqring_wait with no disk progress, matching the observed idle-GPU warmup hang. Depth 4 avoids the pressure window often enough, and the all-tensors-paged 64 MB slot mode has much lower request churn, so it does not trip the same accounting skew.

Fix: track prepared-but-not-yet-submitted req_ids separately from kernel-owned pending requests. flush now retries interrupted submit, drains ready CQEs into the existing demux buffer on -EBUSY/-EAGAIN/zero progress, and only reports ErrorNoSubmit completions for req_ids that cannot be submitted. reap_raw also flushes any prepared SQEs before blocking and rechecks pending so it never calls io_uring_wait_cqe for requests that have already been converted to terminal buffered results. The same contract is applied to the P2P ring.

Regression: added a depth-1 batched FileIOLayer targeted-wait test to exercise chunked submit plus wait_for_req routing without model load or inference. Local build lacks liburing, so the test runs through SyncPread fallback here; the io_uring code path compiled in the target but still needs validation on the GPU/liburing host.

Human validation required: rerun resident-dense warmup/decode with WP_RESIDENT_DENSE=1 and WP_IOURING_DEPTH/WP_PREFETCH_DEPTH at 8 and 16. The server should become ready and show continued page-in/eviction progress instead of blocking in io_cqring_wait.

Assisted-by: Codex

Author: kmbandy

src/weight-pager/wp-file-io-p2p.cpp

@@ -5,6 +5,7 @@
 #include <cerrno>
 #include <cstdlib>
 #include <cstring>
+#include <deque>
 #include <memory>
 #include <vector>
 #include <unistd.h>
@@ -103,7 +104,7 @@ class IoUringP2PFileIOLayer : public FileIOLayer {
 
         struct io_uring_sqe * sqe = io_uring_get_sqe(&ring_);
         if (sqe == nullptr) {
-            io_uring_submit(&ring_);
+            flush_submissions_();
             sqe = io_uring_get_sqe(&ring_);
         }
         if (sqe == nullptr) {
@@ -114,18 +115,19 @@ class IoUringP2PFileIOLayer : public FileIOLayer {
         io_uring_prep_read(sqe, fd_idx, mapped_dst, (unsigned) size, (off_t) offset);
         sqe->flags    |= IOSQE_FIXED_FILE;
         sqe->user_data = req_id;
+        pending_submit_.push_back(req_id);
         ++pending_;
         return true;
     }
 
     void flush() override {
+        if (ring_ok_ && !pending_submit_.empty()) {
+            flush_submissions_();
+        }
         if (!p2p_enabled_ && pending_ == 0) {
             host_->flush();
             return;
         }
-        if (ring_ok_ && pending_ > 0) {
-            io_uring_submit(&ring_);
-        }
     }
 
     // Reap one raw completion for the FileIOLayer base demux. Drains the P2P
@@ -136,6 +138,13 @@ class IoUringP2PFileIOLayer : public FileIOLayer {
     // buffering happen in the base.
     bool reap_raw_(int timeout_ms, IoResult & out) override {
         if (ring_ok_ && pending_ > 0) {
+            if (!pending_submit_.empty()) {
+                flush_submissions_();
+            }
+            if (pending_ == 0) {
+                return false;
+            }
+
             struct io_uring_cqe * cqe = nullptr;
             int ret = 0;
             if (timeout_ms < 0) {
@@ -241,6 +250,90 @@ class IoUringP2PFileIOLayer : public FileIOLayer {
     IoUringP2PFileIOLayer(std::vector<int> fds, std::unique_ptr<FileIOLayer> host)
         : fds_(std::move(fds)), host_(std::move(host)) {}
 
+    void synthesize_submit_failures_(int err) {
+        while (!pending_submit_.empty()) {
+            const uint64_t req_id = pending_submit_.front();
+            pending_submit_.pop_front();
+
+            IoResult r;
+            r.req_id     = req_id;
+            r.status     = IoStatus::ErrorNoSubmit;
+            r.bytes_read = err;
+            ready_[req_id] = r;
+            --pending_;
+        }
+    }
+
+    bool flush_submissions_() {
+        if (!ring_ok_ || pending_submit_.empty()) {
+            return true;
+        }
+
+        while (!pending_submit_.empty()) {
+            int ret = 0;
+            do {
+                ret = io_uring_submit(&ring_);
+            } while (ret == -EINTR);
+
+            if (ret == -EBUSY || ret == -EAGAIN || ret == 0) {
+                bool drained = false;
+                while (reap_ready_cqe_()) {
+                    drained = true;
+                }
+                if (drained) {
+                    continue;
+                }
+            }
+            if (ret < 0) {
+                switch_to_host_errno_("io_uring submit failed", -ret);
+                synthesize_submit_failures_(ret);
+                return false;
+            }
+            if (ret == 0) {
+                switch_to_host_("io_uring submit made no progress", 0);
+                synthesize_submit_failures_(-EAGAIN);
+                return false;
+            }
+
+            for (int i = 0; i < ret && !pending_submit_.empty(); ++i) {
+                pending_submit_.pop_front();
+            }
+        }
+        return true;
+    }
+
+    bool reap_ready_cqe_() {
+        if (pending_ == 0) {
+            return false;
+        }
+
+        struct io_uring_cqe * cqe = nullptr;
+        int ret = io_uring_peek_cqe(&ring_, &cqe);
+        if (ret == -EAGAIN || ret == -EINTR || cqe == nullptr) {
+            return false;
+        }
+        if (ret < 0) {
+            return false;
+        }
+
+        IoResult r;
+        r.req_id = cqe->user_data;
+        const int res = cqe->res;
+        io_uring_cqe_seen(&ring_, cqe);
+        --pending_;
+
+        if (res < 0) {
+            r.status     = IoStatus::ErrorIo;
+            r.bytes_read = res;
+            switch_to_host_errno_("read failed", -res);
+        } else {
+            r.status     = IoStatus::Ok;
+            r.bytes_read = res;
+        }
+        ready_[r.req_id] = r;
+        return true;
+    }
+
     bool init_(int queue_depth, const FileIOP2PConfig & cfg) {
         int rt_version = 0;
         if (hipRuntimeGetVersion(&rt_version) == hipSuccess) {
@@ -386,6 +479,7 @@ class IoUringP2PFileIOLayer : public FileIOLayer {
     bool ring_ok_ = false;
     bool files_registered_ = false;
     int pending_ = 0;
+    std::deque<uint64_t> pending_submit_;
     struct io_uring ring_ {};
     void * libhsa_ = nullptr;
     HsaExportDmaBufFn hsa_export_ = nullptr;

src/weight-pager/wp-file-io.cpp

@@ -233,21 +233,21 @@ class IoUringAsyncFileIO : public FileIOLayer {
         struct io_uring_sqe * sqe = io_uring_get_sqe(&ring_);
         if (sqe == nullptr) {
             // Ring full: flush and retry once.
-            io_uring_submit(&ring_);
+            flush_submissions_();
             sqe = io_uring_get_sqe(&ring_);
             if (sqe == nullptr) return false;
         }
         // Use registered-file index (faster than passing raw fds).
         io_uring_prep_read(sqe, fd_idx, dst, (unsigned) size, (off_t) offset);
         sqe->flags     |= IOSQE_FIXED_FILE;
         sqe->user_data  = req_id;
+        pending_submit_.push_back(req_id);
         ++pending_;
         return true;
     }
 
     void flush() override {
-        if (!ring_ok_ || pending_ == 0) return;
-        io_uring_submit(&ring_);
+        flush_submissions_();
     }
 
     // Reap one completion from the ring. The demux/routing (buffering foreign
@@ -265,6 +265,12 @@ class IoUringAsyncFileIO : public FileIOLayer {
         if (pending_ == 0) {
             return false;  // nothing in flight
         }
+        if (!pending_submit_.empty()) {
+            flush_submissions_();
+        }
+        if (pending_ == 0) {
+            return false;
+        }
 
         struct io_uring_cqe * cqe = nullptr;
         int                   ret = 0;
@@ -347,28 +353,113 @@ class IoUringAsyncFileIO : public FileIOLayer {
             struct io_uring_sqe * sqe = io_uring_get_sqe(&ring_);
             if (sqe == nullptr) {
                 // SQ ring full mid-batch — flush what we've prepped, then retry once.
-                io_uring_submit(&ring_);
+                flush_submissions_();
                 sqe = io_uring_get_sqe(&ring_);
                 if (sqe == nullptr) break;
             }
             io_uring_prep_read(sqe, r.fd_idx, r.dst, (unsigned) r.size, (off_t) r.offset);
             sqe->flags     |= IOSQE_FIXED_FILE;
             sqe->user_data  = r.req_id;
+            pending_submit_.push_back(r.req_id);
             ++pending_;
             ++n_queued;
         }
         // Pass 2: single io_uring_submit for the batch. This is the MAD-235
         // win — one syscall covers N expert-prefetches for the MoE layer
         // instead of N separate submits.
         if (n_queued > 0) {
-            io_uring_submit(&ring_);
+            flush_submissions_();
         }
         return n_queued;
     }
 
 private:
     explicit IoUringAsyncFileIO(std::vector<int> fds) : fds_(std::move(fds)) {}
 
+    void synthesize_submit_failures_(int err) {
+        while (!pending_submit_.empty()) {
+            const uint64_t req_id = pending_submit_.front();
+            pending_submit_.pop_front();
+
+            IoResult r;
+            r.req_id     = req_id;
+            r.status     = IoStatus::ErrorNoSubmit;
+            r.bytes_read = err;
+            ready_[req_id] = r;
+            --pending_;
+        }
+    }
+
+    bool flush_submissions_() {
+        if (!ring_ok_ || pending_submit_.empty()) {
+            return true;
+        }
+
+        while (!pending_submit_.empty()) {
+            int ret = 0;
+            do {
+                ret = io_uring_submit(&ring_);
+            } while (ret == -EINTR);
+
+            if (ret == -EBUSY || ret == -EAGAIN || ret == 0) {
+                bool drained = false;
+                while (reap_ready_cqe_()) {
+                    drained = true;
+                }
+                if (drained) {
+                    continue;
+                }
+            }
+            if (ret < 0) {
+                LLAMA_LOG_WARN("wp::IoUringAsyncFileIO: io_uring_submit failed: %s\n",
+                               strerror(-ret));
+                synthesize_submit_failures_(ret);
+                return false;
+            }
+            if (ret == 0) {
+                LLAMA_LOG_WARN("wp::IoUringAsyncFileIO: io_uring_submit made no progress with %zu SQEs pending\n",
+                               pending_submit_.size());
+                synthesize_submit_failures_(-EAGAIN);
+                return false;
+            }
+
+            for (int i = 0; i < ret && !pending_submit_.empty(); ++i) {
+                pending_submit_.pop_front();
+            }
+        }
+        return true;
+    }
+
+    bool reap_ready_cqe_() {
+        if (pending_ == 0) {
+            return false;
+        }
+
+        struct io_uring_cqe * cqe = nullptr;
+        int ret = io_uring_peek_cqe(&ring_, &cqe);
+        if (ret == -EAGAIN || ret == -EINTR || cqe == nullptr) {
+            return false;
+        }
+        if (ret < 0) {
+            return false;
+        }
+
+        IoResult r;
+        r.req_id = cqe->user_data;
+        const int res = cqe->res;
+        if (res < 0) {
+            r.status     = IoStatus::ErrorIo;
+            r.bytes_read = res;
+        } else {
+            r.status     = IoStatus::Ok;
+            r.bytes_read = res;
+        }
+        io_uring_cqe_seen(&ring_, cqe);
+        --pending_;
+        ready_[r.req_id] = r;
+        return true;
+    }
+
     bool init_(int queue_depth) {
         int ret = io_uring_queue_init(queue_depth, &ring_, 0);
         if (ret < 0) {
@@ -397,6 +488,7 @@ class IoUringAsyncFileIO : public FileIOLayer {
     }
 
     std::vector<int>   fds_;
+    std::deque<uint64_t> pending_submit_;
     bool               ring_ok_          = false;
     bool               files_registered_ = false;
     int                pending_          = 0;

tests/test-weight-pager.cpp

@@ -1383,6 +1383,52 @@ static int test_file_io_submit_batch_partial_failure() {
     return fails;
 }
 
+static int test_file_io_submit_batch_depth_one_targeted_waits() {
+    int fails = 0;
+
+    char path[] = "/tmp/wp-test-batch-depth1-XXXXXX";
+    int fd = mkstemp(path);
+    if (fd < 0) {
+        std::fprintf(stderr, "  FAIL: %s: mkstemp failed: %s\n", __func__, std::strerror(errno));
+        return 1;
+    }
+
+    constexpr size_t N = 16384;
+    std::vector<uint8_t> pattern(N);
+    for (size_t i = 0; i < N; ++i) pattern[i] = (uint8_t) ((i * 19 + 11) & 0xff);
+    ssize_t w = write(fd, pattern.data(), N);
+    EXPECT_EQ_INT((size_t) w, N, "wrote pattern");
+
+    std::vector<int> fds = { fd };
+    auto layer = wp::create_file_io(std::move(fds), /*prefer_async=*/true, 1);
+    EXPECT(layer != nullptr, "create_file_io non-null");
+    if (!layer) { unlink(path); return fails; }
+
+    std::vector<std::vector<uint8_t>> dst(8, std::vector<uint8_t>(1024));
+    std::vector<wp::FileIOBatchRequest> reqs;
+    reqs.reserve(dst.size());
+    for (size_t i = 0; i < dst.size(); ++i) {
+        reqs.push_back({ 1000 + i, 0, (uint64_t) (i * 1536), 1024, dst[i].data() });
+    }
+
+    int n_ok = layer->submit_batch(reqs);
+    EXPECT_EQ_INT(n_ok, (int) reqs.size(), "all depth-1 batch entries accepted");
+
+    for (size_t i = 0; i < reqs.size(); ++i) {
+        wp::IoResult r = layer->wait_for_req(reqs[i].req_id, /*timeout_ms=*/5000);
+        EXPECT(r.status == wp::IoStatus::Ok, "targeted wait returns Ok");
+        EXPECT_EQ_INT(r.req_id, reqs[i].req_id, "targeted wait req_id");
+        EXPECT_EQ_INT(r.bytes_read, 1024, "targeted wait bytes");
+        EXPECT(std::memcmp(dst[i].data(), pattern.data() + i * 1536, 1024) == 0,
+               "targeted wait content");
+    }
+    EXPECT_EQ_INT(layer->pending(), 0, "no pending after targeted waits");
+
+    layer.reset();
+    unlink(path);
+    return fails;
+}
+
 // ---------------------------------------------------------------------------
 // Completion demux — targeted waits must never drop a sibling's completion
 // ---------------------------------------------------------------------------
@@ -1670,6 +1716,7 @@ int main() {
         { "file_io_advise_prefetch",  test_file_io_advise_prefetch  },
         { "file_io_submit_batch",            test_file_io_submit_batch            },
         { "file_io_submit_batch_partial",    test_file_io_submit_batch_partial_failure },
+        { "file_io_submit_batch_depth_one_targeted_waits", test_file_io_submit_batch_depth_one_targeted_waits },
         { "file_io_demux_no_cross_drain",    test_file_io_demux_no_cross_drain    },
         { "compute_advise_ranges",    test_compute_advise_ranges    },
         { "is_uma_archname",          test_is_uma_archname          },