diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index cde0755e8..1ead42647 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -421,6 +421,214 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return ret; } +#if (defined(__linux) || defined(__linux__)) +int LinuxAlignedFileReader::submit(PendingBatch &batch, + std::vector &read_reqs, + IOContext &ctx) { + batch.n_submitted = 0; + batch.n_reaped = 0; + batch.used_pread = false; + batch.cbs.clear(); + batch.cb_ptrs.clear(); + + if (this->file_desc == -1) { + LOG_ERROR("submit: invalid file descriptor"); + return IndexError_Runtime; + } + + if (read_reqs.empty()) { + return 0; + } + + // If no async I/O backend is available, use synchronous pread. + if (ailego::IOBackend::Instance().is_pread()) { + int pread_ret = execute_io_pread(this->file_desc, read_reqs); + if (pread_ret != 0) { + return pread_ret; + } + batch.used_pread = true; + batch.n_submitted = (uint32_t)read_reqs.size(); + return 0; + } + + uint32_t n_ops = (uint32_t)read_reqs.size(); + batch.cbs.resize(n_ops); + batch.cb_ptrs.resize(n_ops); + + for (uint32_t j = 0; j < n_ops; j++) { + io_prep_pread(&batch.cbs[j], this->file_desc, read_reqs[j].buf, + read_reqs[j].len, read_reqs[j].offset); + batch.cbs[j].data = (void *)(uintptr_t)j; + batch.cb_ptrs[j] = &batch.cbs[j]; + } + + int ret = LibAioLoader::Instance().io_submit(ctx, (int64_t)n_ops, + batch.cb_ptrs.data()); + if (ret == (int)n_ops) { + batch.n_submitted = n_ops; + return 0; + } + + // Partial submission: a positive return value means exactly that prefix is + // now in flight and must never be submitted again. Keep submitting the + // remainder; -EAGAIN/-EINTR are transient and worth a bounded retry. + constexpr size_t kMaxSubmitRetries = 8; + uint32_t submitted = (ret > 0 && ret < (int)n_ops) ? (uint32_t)ret : 0; + size_t n_tries = 0; + bool submission_ok = (submitted > 0) || ret == -EAGAIN || ret == -EINTR; + while (submission_ok && submitted < n_ops) { + uint32_t remaining = n_ops - submitted; + ret = LibAioLoader::Instance().io_submit(ctx, (int64_t)remaining, + batch.cb_ptrs.data() + submitted); + if (ret > 0 && (uint32_t)ret <= remaining) { + submitted += (uint32_t)ret; + n_tries = 0; + continue; + } + if ((ret == -EAGAIN || ret == -EINTR) && n_tries < kMaxSubmitRetries) { + n_tries++; + continue; + } + submission_ok = false; + } + + if (submission_ok) { + batch.n_submitted = n_ops; + return 0; + } + + LOG_WARN( + "submit: io_submit stopped after %u/%u requests; returned: %d. " + "falling back to pread after draining submitted AIO", + submitted, n_ops, ret); + + // Drain every request already in flight before any synchronous read can + // reuse its destination buffer, and before batch.cbs may be reused; the + // kernel keeps writing through those iocbs until their events are reaped. + std::vector evts(submitted); + uint32_t drained = 0; + while (drained < submitted) { + uint32_t remaining = submitted - drained; + ret = LibAioLoader::Instance().io_getevents(ctx, (int64_t)remaining, + (int64_t)remaining, + evts.data() + drained, nullptr); + if (ret > 0 && (uint32_t)ret <= remaining) { + drained += (uint32_t)ret; + continue; + } + if (ret == -EINTR) { + continue; + } + LOG_ERROR( + "submit: io_getevents failed while draining %u in-flight requests; " + "returned: %d. resetting the AIO context before falling back to pread", + submitted, ret); + if (!reset_aio_context(ctx)) { + // Do not run pread unless io_destroy confirmed that no request can + // still write into these buffers. + return IndexError_Runtime; + } + break; + } + + int pread_ret = execute_io_pread(this->file_desc, read_reqs); + if (pread_ret != 0) { + return pread_ret; + } + batch.used_pread = true; + batch.n_submitted = n_ops; + return 0; +} + +// Quiesce any requests of the batch still in flight before reporting an +// error, so the kernel cannot keep writing into the caller's buffers or +// leave stale completion events for the next batch on this context. +static void quiesce_batch(PendingBatch &batch, IOContext &ctx) { + if (batch.n_reaped < batch.n_submitted && !batch.used_pread) { + if (reset_aio_context(ctx)) { + batch.n_reaped = batch.n_submitted; + } + } +} + +int LinuxAlignedFileReader::get_completed( + PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices) { + completed_indices.clear(); + + if (batch.n_reaped >= batch.n_submitted) { + return 0; + } + + if (batch.used_pread) { + for (uint32_t i = batch.n_reaped; i < batch.n_submitted; i++) { + completed_indices.push_back(i); + } + batch.n_reaped = batch.n_submitted; + return (int)completed_indices.size(); + } + + uint32_t n_remaining = batch.n_submitted - batch.n_reaped; + int min_req = std::min((int)n_remaining, min_completed); + if (min_req < 1) min_req = 1; + + std::vector evts(n_remaining); + int ret; + do { + // Once requests are in flight, EINTR must be retried: returning here + // would leave them unquiesced, free to overwrite the caller's buffers + // or leak completion events into the next batch. + ret = LibAioLoader::Instance().io_getevents( + ctx, (int64_t)min_req, (int64_t)n_remaining, evts.data(), nullptr); + } while (ret == -EINTR); + if (ret < 0) { + LOG_ERROR("get_completed: io_getevents failed, ret=%d, %s", ret, + ::strerror(-ret)); + quiesce_batch(batch, ctx); + return IndexError_Runtime; + } + + for (int i = 0; i < ret; i++) { + uint32_t idx = (uint32_t)(uintptr_t)evts[i].data; + if (idx >= batch.n_submitted) { + LOG_ERROR("get_completed: completion referenced an unknown request %u", + idx); + batch.n_reaped += (uint32_t)ret; + quiesce_batch(batch, ctx); + return IndexError_Runtime; + } + int64_t res = (int64_t)evts[i].res; + int64_t res2 = (int64_t)evts[i].res2; + int64_t expected = (int64_t)batch.cbs[idx].u.c.nbytes; + if (res != expected || res2 != 0) { + // The async read failed, so the destination buffer content is + // undefined. Degrade to a synchronous pread for this request before + // handing the buffer to the caller. + LOG_WARN( + "get_completed: read %u failed: res=%ld, res2=%ld, expected=%ld; " + "retrying with pread", + idx, (long)res, (long)res2, (long)expected); + ssize_t bytes_read = + ::pread(this->file_desc, batch.cbs[idx].u.c.buf, + batch.cbs[idx].u.c.nbytes, batch.cbs[idx].u.c.offset); + if (bytes_read != (ssize_t)expected) { + LOG_ERROR( + "get_completed: pread retry for read %u failed; got=%zd, " + "expected=%ld, errno=%d, %s", + idx, bytes_read, (long)expected, errno, ::strerror(errno)); + batch.n_reaped += (uint32_t)ret; + quiesce_batch(batch, ctx); + return IndexError_Runtime; + } + } + completed_indices.push_back(idx); + } + + batch.n_reaped += (uint32_t)ret; + return ret; +} +#endif + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index a1cb7c91a..a8e1354fd 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -58,6 +58,16 @@ struct AlignedRead { } }; +#if (defined(__linux) || defined(__linux__)) +struct PendingBatch { + std::vector cbs; + std::vector cb_ptrs; + uint32_t n_submitted{0}; + uint32_t n_reaped{0}; + bool used_pread{false}; +}; +#endif + class AlignedFileReader { protected: std::map ctx_map; @@ -77,6 +87,15 @@ class AlignedFileReader { virtual int read(std::vector &read_reqs, IOContext &ctx, bool async = false) = 0; + +#if (defined(__linux) || defined(__linux__)) + virtual int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx) = 0; + + virtual int get_completed(PendingBatch &batch, IOContext &ctx, + int min_completed, + std::vector &completed_indices) = 0; +#endif }; class LinuxAlignedFileReader : public AlignedFileReader { @@ -101,6 +120,14 @@ class LinuxAlignedFileReader : public AlignedFileReader { int read(std::vector &read_reqs, IOContext &ctx, bool async = false); + +#if (defined(__linux) || defined(__linux__)) + int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx); + + int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices); +#endif }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..987a7da76 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -80,7 +80,7 @@ int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { sector_num_per_node_ = DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); - if (beam_width_ > sector_num_per_node_ * DiskAnnUtil::kMaxSectorReadNum) { + if (beam_width_ * sector_num_per_node_ > DiskAnnUtil::kMaxSectorReadNum) { LOG_ERROR("Beamwidth can not be higher than kMaxSectorReadNum"); return IndexError_InvalidArgument; @@ -791,18 +791,29 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { uint32_t num_ios = 0; + // Cap beam width so one batch of frontier reads never exceeds the sector + // buffer capacity (kMaxSectorReadNum sectors), as each node occupies + // sector_num_per_node sectors. + uint32_t max_beam_width = + std::max(1u, static_cast(DiskAnnUtil::kMaxSectorReadNum / + sector_num_per_node)); + uint32_t effective_beam_width = std::min( + std::max(8u, std::min(ctx->list_size() / 5, 32u)), max_beam_width); + std::vector frontier; - frontier.reserve(2 * beam_width_); + frontier.reserve(2 * effective_beam_width); std::vector> frontier_neighbors; - frontier_neighbors.reserve(2 * beam_width_); + frontier_neighbors.reserve(2 * effective_beam_width); std::vector frontier_read_reqs; - frontier_read_reqs.reserve(2 * beam_width_); + frontier_read_reqs.reserve(2 * effective_beam_width); std::vector> cached_neighbors; - cached_neighbors.reserve(2 * beam_width_); + cached_neighbors.reserve(2 * effective_beam_width); + + PendingBatch pending; while (candidates.has_unexpanded_node() && num_ios < io_limit_) { frontier.clear(); @@ -813,8 +824,9 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { uint64_t sector_buffer_idx = 0; uint32_t num_seen = 0; - while (candidates.has_unexpanded_node() && frontier.size() < beam_width_ && - num_seen < beam_width_) { + while (candidates.has_unexpanded_node() && + frontier.size() < effective_beam_width && + num_seen < effective_beam_width) { auto neighbor = candidates.closest_unexpanded(); num_seen++; @@ -857,11 +869,10 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { } io_timer.reset(); - - int read_ret = reader_->read(frontier_read_reqs, io_ctx); + int submit_ret = reader_->submit(pending, frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); - if (read_ret != 0) { - LOG_ERROR("cached_beam_search: reader_->read failed, ret=%d", read_ret); + if (submit_ret != 0) { + LOG_ERROR("cached_beam_search: submit failed, ret=%d", submit_ret); ctx->set_error(true); return IndexError_Runtime; } @@ -895,55 +906,73 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { for (uint64_t m = 0; m < neighbor_num; ++m) { diskann_id_t id = node_neighbors[m]; - visit_filter.set_visited(id); - - Neighbor nn(id, distances[m]); - candidates.insert(nn); + if (!visit_filter.visited(id)) { + visit_filter.set_visited(id); + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } } } - for (auto &frontier_neighbor : frontier_neighbors) { - uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( - node_per_sector_, max_node_size_, frontier_neighbor.second, - frontier_neighbor.first); - uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( - node_disk_buf, meta_.element_size()); - uint32_t neighbor_num = *node_buf; - - void *node_fp_coords = node_disk_buf; - - float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords); + if (!frontier.empty()) { + std::vector completed; + while (pending.n_reaped < pending.n_submitted) { + completed.clear(); + io_timer.reset(); + int n = reader_->get_completed(pending, io_ctx, 1, completed); + stats.io_us += io_timer.micro_seconds(); + if (n < 0) { + LOG_ERROR("cached_beam_search: get_completed failed, ret=%d", n); + ctx->set_error(true); + return IndexError_Runtime; + } - if (!ctx->filter().is_valid() || - !ctx->filter()(get_key(frontier_neighbor.first))) { - topk_heap.emplace( - frontier_neighbor.first, - VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords))); - } + for (uint32_t idx : completed) { + auto &frontier_neighbor = frontier_neighbors[idx]; + uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( + node_per_sector_, max_node_size_, frontier_neighbor.second, + frontier_neighbor.first); + uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( + node_disk_buf, meta_.element_size()); + uint32_t neighbor_num = *node_buf; - diskann_id_t *node_neighbors = - reinterpret_cast(node_buf + 1); + void *node_fp_coords = node_disk_buf; - cpu_timer.reset(); - std::vector distances(neighbor_num); - pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, - ctx->pq_table_dist_buffer(), - ctx->pq_coord_buffer(), distances.data()); + float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords); - stats.dist_num += neighbor_num; - stats.cpu_us += cpu_timer.micro_seconds(); + if (!ctx->filter().is_valid() || + !ctx->filter()(get_key(frontier_neighbor.first))) { + topk_heap.emplace(frontier_neighbor.first, + VectorInfo(cur_expanded_dist, + make_vector_copy(node_fp_coords))); + } - cpu_timer.reset(); - for (uint64_t m = 0; m < neighbor_num; ++m) { - diskann_id_t id = node_neighbors[m]; - visit_filter.set_visited(id); - stats.dist_num++; + diskann_id_t *node_neighbors = + reinterpret_cast(node_buf + 1); + + cpu_timer.reset(); + std::vector distances(neighbor_num); + pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, + ctx->pq_table_dist_buffer(), + ctx->pq_coord_buffer(), distances.data()); + + stats.dist_num += neighbor_num; + stats.cpu_us += cpu_timer.micro_seconds(); + + cpu_timer.reset(); + for (uint64_t m = 0; m < neighbor_num; ++m) { + diskann_id_t id = node_neighbors[m]; + if (!visit_filter.visited(id)) { + visit_filter.set_visited(id); + stats.dist_num++; + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } + } - Neighbor nn(id, distances[m]); - candidates.insert(nn); + stats.cpu_us += cpu_timer.micro_seconds(); + } } - - stats.cpu_us += cpu_timer.micro_seconds(); } }